目录文章目录目录Decode-only 推理过程Prefill 和 Decode 阶段KV CacheKV Cache 的基本原理KV Cache 的生成流程KV Cache 的数学原理KV Cache 的容量计算KV Cache 的计算量计算PD 分离PD 分离架构推理性能指标PD 分离差异化配置KV Cache 传输技术 —— MooncakePrefill PoolKV Cache Pool工作流程SGLang 集成Decode-only 推理过程Decode-only Transformer 是一个自回归模型Autoregressive model —— 假定我们输入一个文本模型就会输出一个长度为 N 的回答。那么实际上模型过程中执行了 N 次推理过程且一次推理只输出一个 token。当前轮输出的 token 会与之前输入的 tokens 拼接在一起并作为下一轮的输入 tokens这样不断重复直至遇见 EOS 终止符或生成的 token 数目达到设置的 max_new_token 为止停下。Prefill 和 Decode 阶段具体而言自回归模型的推理过程会被划分为 2 个阶段Prefill预填充阶段又称为 Prompt phase提示词处理阶段模型接受到用户输入的 PromptIs tomato a fruit?根据输入 TokensIs, tomato, a, fruit, ? 生成第一个输出 TokenYes。Decode增量解码阶段又称为 Token-generation phasetoken 生成阶段从生成第一个输出 TokenYes 之后开始把 Prompt 以及已生成的输出 tokens 组成新的模型输入采用自回归方式一次生成一个 Token直到结束。PrefillDecode实际作用一次性计算 Prompt 全序列的注意力得分、生成第一个输出 token、生成初始 KV cache、生成一个新标记作为 Decode 阶段的初始输入。收到 Prefill 的全部输出后进入自回归过程开始逐轮/逐一生成输出 token。触发时机用户输入 Prompt 后触发Prefill 生成第一个输出 token 后触发输入内容PromptPrompt 和上一次的输出 tokens自回归执行过程执行一次前向计算执行 N-1 次前向计算输入序列长度为 N输出内容输出第一个 token、新标记、Prompt 全序列的 KV cache。每一轮输出一个 token和这个 token 的 KV cache。计算类型计算密集型任务存在大量 GEMMGEneral Matrix-Matrix multiply一般矩阵乘法操作访存密集型任务存在大量的 KV cache 读写操作并行特性Prompt 全序列进行 MHA 并行计算并行度高自回归存在串行顺序性可见自回归模型的生成模式是存在 2 个阶段的原因其中的 KV cache 是主要的优化方向。KV Cache存在 KV CacheKV 矩阵缓存的原因是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]输入 “新”输出 “年。将 “年” 拼接到 “新” 的后面作为新的输入即本次推理的输入为 “新年”预测得到 “快”。将 “快” 拼接到 “新年” 的后面作为新的输入即本次推理的输入为 “新年快”预测得到 “乐”。上述例子可见Decode 每轮生成输出 token 时都可以复用前面轮次 token 的 KV 矩阵。如果没有 KV cache 的话那么每轮都要重新算一次所有历史 token 的 KV 矩阵复杂度为 O(n^2)。这显然是没必要的。并且实际上一次前向计算中的多个位置都存在 KV 的冗余计算如下图所示。embedding 操作KV 矩阵生成操作Q×K^T 操作softmax 的注意力得分与 V 相乘操作KV Cache 的生成流程输入 prompt “新年快”。此时会把 “新年快” 的 KV 矩阵计算出来存储在 KV Cache 中。然后输出 token “乐”。下一轮只需要计算 “乐” 的 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 × N2代表 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 的层数假设总上下文 100K60 层8 个头128 的嵌入维度使用 BF16 存储则 KV Cache 大小为可见KV Cache 由输入输出序列决定和模型权重参数没有关系。随着推理的 batch size、sequencxe length 变大变长那么 KV Cache 占用的存储空间很可能超过模型本身。所以变长数据的 KV cache 存储是一个比较大的难题。KV Cache 的计算量计算由于 KV Cache 的存在Mask Attention 计算的时候从矩阵乘法降为矩阵乘向量计算量的比较如下No KV Cache24 b s h 2 4 b s 2 h 24bsh^2 4bs^2h24bsh24bs2hKV Cache24 b h 2 4 b s h 24bh^2 4bsh24bh24bsh可见对于单次运算KV cache 模型下的计算量减少了 s 倍。PD 分离具体而言Prefill 和 Decode 的任务特性存在较大的差异Preffill 强关联 TTFT、Decode 强关联 TPOT。因此将 P、D 分开部署在不同的设备上一方面消除了二者之间的干扰。另一方面也更有利于差异优化使得两个阶段都能更快地达到各自的 SLOService 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 分离已经成为了主流。阶段 1No KV Cache 时期阶段 2KV Cache 时期阶段 3PD 分离时期如上图所示PD 分离的处理流程主要包括 3 个阶段Prefill当一个请求进入系统后首先发送到一个 Prefill Instance 执行 Prefill 运算得到第一个 token 以及 prompt 对应的 KV Cache。Migrate将 Prefill 处理后的 token、KV Cache 迁移至 Decode InstanceDecode Decode Instance 对迁移的请求执行 Decode 运算。如上图所示KV Cache 的存储和传输就成为了 PD 分离架构的关键之一。Prefill 初始化 KV Cache 存储Decode 检索和更新 KV Cache 存储。这个 KV Cache 存储可能是 P 或 D 设备上的存储也可能是一个外挂的第三方专用存储设备。推理性能指标评价模型推理过程的性能指标主要有TTFTTime to first token首 token 时延单位时间是用户发出请求后模型输出第一个 token 所需的时间它主要受 Prefill 阶段计算和缓存加载影响决定了响应的 “开口速度”TPOTTime per output token每 token 生成时延单位时间是模型在开始生成后平均每个 token 的输出耗时它反映了 Decode 阶段的效率决定了生成过程的 “语速流畅度”。TPS输出 token 吞吐量单位tok/sPD 分离差异化配置Prefill 和 Decode 的特性很不同所以参数配置上也有很多的不同。例如批策略量化类型TP 大小PP 大小任务超时时间重试时间显存申请策略RDMA 软硬件队列大小是否可回退执行等。P 阶段D 阶段资源分配选择计算能力强的卡型选择显存大的卡型一次前向计算处理的 token 数8k-16k256配置比例13量化方案部署 FP16 / W8A8 的模型文件来获得相对不错的 Prefill 性能。自由地选择 W4A16 / W4A8 的方案来获得整体的更优性能。注 1为了充分发挥 Decode 的能力应该配置更多的 Prefill 数量和更少的 Decode 数量以此来提供更多的 token 到 Decode 进行连续批处理。注 2W4A8 表示权重使用 4-bit 量化而激活值保留较高的 16-bit 精度。激活值更接近原始精度通常用于 FP16 或 BF16。KV Cache 传输技术 —— MooncakeMooncake 是一个以 KV Cache 存储为中心的推理引擎组件主要负责 KV Cache 的存储和传输。在 SGLang 推理引擎中Mooncake 作为一种 KV connector backend 实现。上图是 Mooncake 的软件架构图包括KVCache-centric Conductor根据当前的 KV Cache 分布和工作负载分派请求。Prefill Pool处理用户输入的 Prompt目标在 Prefill 阶段也尽可能多的重用 KV Cache以避免冗余计算进而优化 TTFT。Decoding Pool自回归式的流式输出影响 TBTtime-between-tokens。KV Cache Pool利用每台 GPU 服务器上的主存、SSD 磁盘等存储空间组成一个 KVCache Pool 来进行全局的 Prefix Cache。进而全局 Prefix Cache 通过全局的调度能够大幅度提升复用率从而提升总吞吐。Prefill PoolPrefill Pool 实现了以下关键技术CPPChunked pipeline parallelism通过 Chunk 分块流水线并行机制来扩展单个请求的处理跨多个节点。并且以流水线的模式结合 KV Cache 的异构传输来支持 Prefill 和 Decode 流程的 overlap。Layer-Wise Prefill分层完成 Prefill。每个 MMHA 层的注意力计算开始之前模型会等待该层的 KV Cache 的异步加载完成并触发下一层的异步 KV Cache 加载。在注意力计算完成后会启动该层 KV Cache 的异步存储。Multi-Node Prefill分块流水线并行对于每个请求其输入标记被分成不同的 Chunk 块每个块不超过预填充块。同一请求的不同块可以由不同节点同时处理从而实现处理的并行化并减少 TTFT。KV Cache Pool使用 KV Cache Pool 后在 P 和 D 之间增加了一个中间存储Prefill 节点先将 KV Cache 写到中间存储Decode 节点从中间存储读。KV Cache Pool 实现了以下技术Chunk 块管理KV Cache 以 Chunk 为单位进行组织每个 Chunk 块都附有一个哈希值该哈希值由其自身的哈希值和 Prefix 决定。支持多会话的共享缓存全局范围内避免重复存储和计算。支持分层存储管理高频 KV Cache 驻留 CPU 主存低频数据下沉至 SSD 磁盘。支持 RDMA 传输KV Cache 在 CPU 和 GPU 之间的传输由一个支持 GPUDirect / RDMA 的 Messenger 模块处理。缓存驱逐算法如 LRU最近最少使用、LFU最不频繁使用或基于请求特征的缓存淘汰算法算会会将 KV Cache 转入 CPU 内存。工作流程对于每个新请求Conductor 根据请求的 seq_len 和 prefix_len 估计相应的执行时间。因实例而异将请求的估计等待时间相加以获取该实例上的预计 TTFT。Conductor 将请求分配给 TTFT 最短的实例并相应地更新该实例的缓存和队列时间。至此Conductor 选择了一对 Prefill 节点和一个 Decode 节点。Conductor 根据 Prefix Cache 的 Chunk ID 将 Prefix Cache 从远程 CPU 主存加载到 GPU 显存以启动请求KV Cache 重用。Prefill 节点或组使用 Prefix Cache 完成 Prefill 阶段并将新生成的增量 KV Cache 存储回 CPU 主存。如果未命中缓存的输入 token 数量超过了一定阈值prefill_chunk则将 Prefill 阶段分成多个 Chunk 并以流水线方式执行。Messanger 将每个模型层生成的 KV Cache 流式传输到目标 Decode 节点的 CPU 主存此 KV Cache 的传输是异步执行的并与上述增量预填充步骤重叠这样可以减少等待时间。当 Decode 节点的 CPU 主存中接收到所有的 KV Cache 后Mooncake 会以连续批处理的方式加入下一批请求。Conductor 根据 Decode 节点当前的负载预先选择 Decode 节点以保证 TBT SLO 要求。SGLang 集成SGLang 的 KV connector 都有四个角色KVManager是 KV connector 的管理器Prefill 和 Decode 都有。负责初始化内存、BootstrapServer 以及发送 KV Cache。KVBootstrap / MooncakeKVBootstrapServer只用于 Prefill记录 Prefill 和 Decode 交互所用的信息可以有很多种Decode 会请求该信息和 Prefill 进行连接。KVSender专属于 Prefill用于发送 KV Cache。KVReceiver专属于 Decode用于和 Prefill 握手获取 BootstrapServer 的交互信息也用于接收 KV Cache。
