网站开发设计jw100系统设置友情链接有什么作用

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 分离的核心问题

1. 如何传输 KV cache？

   • Pooling mode：在这种模式下，发送方将 KV cache 写入到一个共享的 KV cache 池中，接收方从池中获取所需的 KV cache。这种模式的优点是简单易实现，但可能会因为池的共享访问而导致延迟。
   • P2P mode：在这种模式下，发送方需要知道接收方是谁，并直接建立连接传输 KV cache。这种模式的效果通常比 Pooling mode 更好，但代码设计会更加复杂。

2. 如何从 vLLM 中提取（和注入）KV cache？

   • Connector API：通过设计一个连接器 API，可以在模型前向传播之前接收 KV cache（将 KV cache 注入到 vLLM 的分页内存中），在模型前向传播之后，从 vLLM 的分页内存中提取 KV cache 并发送出去。

3. 何时发送请求到 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 传输和管理的库。