PD分离：优化大语言模型推理效率
在大语言模型的推理过程中，Prefill 和 Decode 是两个关键阶段。随着模型规模的不断扩大，如何高效地处理这两个阶段的计算任务，成为了一个亟待解决的问题。
一、什么是 Prefill 和 Decode？
在大语言模型的推理过程中，Prefill 和 Decode 是两个不同的阶段，它们各自承担着不同的任务。
- Prefill：Prefill 是推理过程的初始阶段，它的主要任务是提前计算好输入文本的 KV cache（键值缓存），并生成第一个 token。这个阶段通常需要处理大量的计算任务，因此速度相对较慢。
- Decode：Decode 是在 Prefill 完成后进行的阶段，它的任务是基于已经生成的 KV cache 不断生成后续的 token。由于 Decode 阶段可以利用已经计算好的 KV cache，因此速度相对较快。
初始逻辑问题
在传统的推理流程中，Prefill 和 Decode 是顺序执行的。然而，这种设计存在一个明显的问题：当一些用户正在处于 Decode 阶段时，如果突然来了一个新的 Prefill 请求，系统会优先处理 Prefill 请求，这会导致正在 Decode 的用户被迫中断。这种中断不仅会影响用户体验，还会降低系统的整体效率。
二、PD 分离的解决方案
为了解决上述问题，研究者们提出了 PD 分离的概念。PD 分离的核心思想是将 Prefill 和 Decode 分别放在两个不同的 GPU 上运行，一个 GPU 专门负责 Prefill，另一个 GPU 专门负责 Decode。这样可以避免 Prefill 和 Decode 之间的冲突，提高系统的整体效率。
PD 分离的核心问题
-
如何传输 KV cache？
- Pooling mode：在这种模式下，发送方将 KV cache 写入到一个共享的 KV cache 池中，接收方从池中获取所需的 KV cache。这种模式的优点是简单易实现，但可能会因为池的共享访问而导致延迟。
- P2P mode：在这种模式下，发送方需要知道接收方是谁，并直接建立连接传输 KV cache。这种模式的效果通常比 Pooling mode 更好，但代码设计会更加复杂。
-
如何从 vLLM 中提取（和注入）KV cache？
- Connector API：通过设计一个连接器 API，可以在模型前向传播之前接收 KV cache（将 KV cache 注入到 vLLM 的分页内存中），在模型前向传播之后，从 vLLM 的分页内存中提取 KV cache 并发送出去。
-
何时发送请求到 P 和 D 节点？
- First P then D：Prefill 先收到请求，处理完成后将结果发送给 Decode。
- First D then P：Decode 先收到请求，如果发现没有 KV cache，则请求 Prefill。
三、其他相关想法
除了 PD 分离，还有一些其他的想法可以进一步优化推理效率：
- Chunked Prefill：让算子可以计算不同的 shape，从而提高 Prefill 阶段的灵活性和效率。
四、相关工具和库
随着 PD 分离概念的提出，一些新的工具和库也应运而生，例如：
- LM Cache：用于高效管理 KV cache 的工具。
- MoonCake：一种支持 PD 分离的框架。
- NIXL：用于优化 KV cache 传输和管理的库。
