24年6月AI公司月之暗面的技术报告“Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving”。
Mooncake 是 Kimi 的服务平台，Kimi 是月之暗面公司提供的一项 LLM 服务。它采用以 K-V Cache 为中心的分解式架构，将预填充和解码 cluster 分开。它还利用 GPU cluster 中未充分利用的 CPU、DRAM 和 SSD 资源来实现 K-V Cache 的分解式缓存。
Mooncake 的核心是以 K-V Cache 为中心的调度程序，在最大化整体有效吞吐量和满足与延迟相关的服务级目标 (SLO) 之间取得平衡。与假设所有请求都将被处理的传统研究不同，Mooncake 面临着高度超载场景带来的挑战。为了缓解这些问题，开发一种基于预测的早期放弃（early rejection）策略。实验表明，Mooncake 在长上下文场景中表现出色。
随着大语言模型 (LLM) 在各种场景中的快速应用 [1、2、3、4]，LLM 服务的工作负载变得非常多样化。这些工作负载在输入/输出长度、频率和分布方面有所不同，最重要的是，需要不同类型的服务级目标 (SLO)。作为模型即服务 (MaaS) 提供商，Kimi [5] 的主要目标之一是解决具有多个复杂约束的优化问题。优化目标是最大化整体有效吞吐量，这直接影响收入，而约束反映了不同级的 SLO。这些 SLO 通常涉及满足与延迟相关的要求，主要是第一个token的时间 (TTFT) 和token之间的时间 (TBT)。
为了实现这一目标，先决条件是充分利用 GPU 集群中可用的各种资源。具体而言，尽管 GPU 服务器目前以高度集成的节点形式提供（例如 DGX/HGX 超级计算机 [6]），但有必要将它们解耦并重组为几个分散的资源池，每个资源池针对不同但协作的目标进行优化。例如，许多研究人员 [7、8、9] 建议将预填充服务器与解码服务器分开，因为 LLM 服务的这两个阶段具有非常不同的计算特性，其中 K-V Cache 会随着请求从预填充服务器转移到解码服务器而发生变化。
基于此，KV Cache 的调度是 LLM 服务调度的核心。为了提高整体吞吐量，通常有两种通用方法：1）尽可能多地重用 KV Cache 减少所需的计算资源；2）最大化每批次的tokens数量以提高模型 FLOPS 利用率 (MFU)。但是，从远程位置重用 KV Cache 会延长 TTFT，而较大的批次会导致更大的 TBT。因此，这两种面向吞吐量优化的利用可能会导致违反与延迟相关的 SLO。
现代大语言模型 (LLM) 基于 Transformer 架构，该架构利用注意机制和多层感知器 (MLP) 来处理输入。流行的基于 Transformer 的模型，例如 GPT [10] 和 LLaMA [11]，采用仅解码器结构。每个推理请求在逻辑上分为两个阶段：预填充阶段和解码阶段。
在预填充阶段，所有输入 token 都并行处理。此阶段生成第一个输出 token，同时存储计算的K-V的中间结果，称为 KV Cache。然后，解码阶段使用此 KV Cache 自回归生成新 token，将计算中的新K-V添加到 KV Cache。除了短请求之外，在预填充阶段同时处理输入 token 的能力，通常会使其计算密集。由于注意网络的计算复杂度随输入长度二次方增加，而 MLP 的复杂度随输入长度线性增加，因此预填充阶段的计算时间通常随输入长度超线性增加，如图左侧所示。
相反，由于自回归生成的限制，解码阶段每批每次只处理一个 token。这使得它受到内存限制，并导致计算时间随批大小亚线性增加，如上图右侧所示。解码阶段广泛使用的优化，是连续批处理 [12, 13]。在每次迭代之前，调度程序都会检查所有请求的状态，将新到达的请求添加到批处理的预填充阶段，同时删除已完成的请求。
由于预填充和解码阶段的特点不同，MaaS 提供商设置不同的指标来衡量其相应的服务级目标 (SLO)。具体而言，预填充阶段主要关注TTFT。另一方面，解码阶段关注的是TBT。
作为 MaaS 提供商，通过满足服务协议定义的 SLO 指标来确保质量保证，至关重要。例如，TTFT/P90 = 4× 这样的指标表示， 90% 的推理请求，其 TTFT 不会大于在相同条件下运行的单个请求四倍，也不会受到干扰。在实际部署中，设置TTFT 和 TBT 固定的 SLO。如果监控检测到未满足的 SLO，会添加推理资源或拒绝一些传入请求。
但是，由于目前 GPU 的供应不稳定，弹性扩展推理 cluster 通常是不可行的。因此，决定拒绝哪些请求成为面向过载的调度核心问题。主要目标是在遵守 SLO 的同时最大化整体吞吐量，在其他研究中，SLO 被称为有效吞吐量 [8, 14]。不同之处在于，只有完全完成执行的请求，才会被计入有效吞吐量的衡量标准。否则，所有之前消耗/生成的tokens都不会被计算在内，相应的资源就会被浪费。换句话说，如果请求无法在 SLO 下完成其完整执行，则应尽早拒绝该请求。实现这一目标不仅需要优化预填充和解码阶段的架构，还需要开发预测短期未来负载的能力。
如图是KV Cache 为中心的 LLM 服务分解架构，名为 Mooncake。对于每个请求，全局调度器 (Conductor) 需要选择一对预填充和解码实例，并按以下步骤调度请求：1) 将尽可能多的可重用 KV Cache 传输到选定的预填充实例；2) 以块/层的形式完成预填充阶段，并将输出 KV Cache 连续流传输到相应的解码实例；3) 加载 KV Cache 并将请求添加到解码实例的连续批处理过程中，生成请求输出。
Mooncake 采用分解式架构，不仅将预填充节点与解码节点分离，还将 GPU 集群的 CPU、DRAM、SSD 和 RDMA 资源分组，以实现分解式 KV Cache。这种分解式缓存可充分利用未充分利用的资源，提供充足的缓存容量和传输带宽，从而无需额外成本即可实现高效的近 GPU 前缀缓存（prefix-caching）。
如下图说明 KV Cache 块的存储和传输逻辑。在 CPU 内存中，KV Cache 以分页块形式存储。根据请求模式，它可以使用缓存驱逐（cache eviction）算法，例如 LRU（Least Recently Used）、LFU（Least Frequently Used）或基于请求特性的算法。这些 KV Cache 块在 CPU 和 GPU 之间的传输由一个单独的（GPUDirect）基于 RDMA 的组件 （即Messenger） 处理。该架构还向外部用户提供上下文缓存 API，以提高 KV Cache 的重用率。
为了调度所有这些分解的组件，Mooncake 在其核心实现了一个名为 Conductor 的全局调度程序。Conductor 负责根据 KV Cache 和工作负载的当前分布调度请求。如果有利于未来推理，它还会复制或交换 KV Cache 的某些块。具体来说，如图所示演示请求的典型工作流程。token化完成后，conductor 选择一对预填充节点和一个解码节点，并启动包含四个步骤的工作流程：
- KV Cache 重用：选定的预填充节点（组）接收一个包含原始输入、可重用前缀缓存块 ID 以及分配给请求的完整缓存块 ID 的请求。它根据前缀缓存块 ID 将前缀缓存从远程 CPU 内存加载到 GPU 内存中以引导请求。如果不存在前缀缓存，则跳过此步骤。此选择平衡三个目标：尽可能多地重用 KVCache、平衡不同预填充节点的工作负载以及保证 TTFT SLO。
- 增量预填充：预填充节点（组）使用前缀缓存完成预填充阶段，并将新生成的步进 KV Cache 存回 CPU 内存。如果未缓存的输入 token 数量超过某个阈值（prefill_chunk），则预填充阶段将拆分为多个块并以流水线方式执行。此阈值的选择是为了充分利用相应 GPU 的计算能力，通常大于 1000 个 token。
3）KV Cache 传输：上述 Messenger 服务部署在每个节点中，管理和传输这些缓存。每个 Messenger 在其各自的推理实例中作为独立进程运行，接收信号以促进高速、跨机器的 KV Cache 传输。此步骤异步执行并与上述步进预填充步骤重叠，将每个模型层生成的 KV Cache 流传输到目标解码节点的 CPU 内存，以减少等待时间。
4）解码：在解码节点的 CPU DRAM 中接收到所有 KV Cache 后，请求以连续批处理的方式加入下一个批处理。 Conductor 根据其当前负载预选择解码节点，确保其不违反 TBT SLO。但是，局部调度程序会对此 SLO 进行双重检查，因为预期负载可能在预填充阶段后发生变化。这种双重检查可能会导致请求放弃，在这种情况下，相应的预填充成本会浪费掉。
分块预填充将输入tokens分成多个块，加入连续批处理过程。这种方法有两个明显的好处：1) 无需分离，所有节点都得到平等对待，使调度更容易；2) 将分块预填充内联到解码批处理中，可以提高解码批处理的计算强度，从而实现更好的 MFU。
MoonCake的做法，仅当请求的预填充可以在不分块的情况下转发且不会损害 TBT SLO 时，才会内联到解码批处理中。做出此决定的主要原因有两个：1) 预填充节点需要不同的跨节点并行设置来处理长上下文。 2) 提供节省 VRAM 的特殊机会。
最近的 LLM 的可用上下文长度正在迅速增加，从 8k 增加到 128K 甚至 1M [16]。通常，对于这种长上下文请求，输入 tokens 可能比输出 tokens 大 10 到 100 倍，因此优化 TTFT 至关重要。由于长上下文预填充中存在大量并行性，因此最好使用多个 8x GPU 节点来并行处理。但是，将张量并行 (TP) 扩展到多个节点需要每层执行两个昂贵的基于 RDMA all-reduce 操作，从而显著减少预填充节点的 MFU。
Mooncake 利用仅解码器 Transformer 的自回归特性，并为长上下文预填充实现分块流水线并行 (CPP)。将预填充cluster 中的每 X 个节点分组为流水线预填充节点组。对于每个请求，其输入 tokens 被划分为块，每个块的长度不超过 prefill_chunk。同一请求的不同块可以由不同的节点同时处理，从而并行化处理并减少 TTFT。
CPP 有两个主要优点：1) 与训练中的流水线并行（PP）类似，它只需要在每个流水线阶段的边界进行跨节点通信，这很容易与计算重叠。这可以实现更好的 MFU 并减少与 KV Cache 传输的网络资源争用。2) 它自然适合短上下文和长上下文，不会给短上下文预填充带来显着的开销，并避免频繁动态调整节点分区。这种基于流水线的加速方法已经在训练系统 [24] 中进行了探索。
除了计算能力之外，有限的 VRAM 大小也是宝贵的资源，目标是最大限度地减少（主要是 KV Cache）对 VRAM 的占用。理论上，如果一个请求的 KV Cache 大小为 S ，处理时间为 T ，则它的占用成本为 S ∗ T 。如果在分块预填充中将请求分块，并且每个块的处理都与其他解码请求内联，则 T 将增加，导致更大的占用成本。
此外，由于预填充是逐层处理的并且受计算限制，因此可以将 KV Cache 的传输和转储与计算重叠，进一步降低其占用成本。在 Mooncake 中，通过启动和等待操作 KV Cache 的加载和存储异步执行。在每一层的注意计算开始之前，模型等待该层的 KV Cache 异步加载完成，并触发下一层异步 KV Cache 加载。在注意计算完成后，启动该层的 KV Cache 异步存储。一旦所有层的计算都完成，该过程将等待所有异步存储操作的完成。传输重叠允许预填充实例的执行时间大致相当于 KVCache 加载时间或标准预填充时间，具体取决于前缀缓存相对于输入长度的比例。KV Cache 存储延迟的实验结果（如图）表明，分层预填充可以有效减少长上下文请求的延迟。
这种重叠有效性的主要优点是，它在预填充调度中忽略可用的 VRAM 大小，只要它可以包含单个请求即可。预填充节点的调度仅考虑 KV Cache 分布和可用的 DRAM 大小。
OpenAI 最近提出了使用批处理 API [25]，这使用户能够以降低 50% 的成本发送异步请求组，但周转时间仅为 24 小时。此服务非常适合处理不需要立即响应的作业。由于这些批处理请求没有严格的 TBT，如果有足够的 VRAM 空间来容纳相应的 KVCache，甚至可以将这些请求的解码阶段内联到预填充处理中，以实现更好的 MFU。
下图是KV Cache为中心的调度算法：
目前大多数关于 LLM 服务的研究都假设所有请求都将被处理，从而优化吞吐量或请求的 TTFT 和 TBT。然而，在实际场景中，处理每个传入请求既不经济也不现实。对于面临快速增长的用户请求量的商业推理服务，cluster推理资源的增长速度远远低于传入请求的增长速度。因此，过载是当前 LLM 服务中的一个常见问题，尤其是在高峰时段。
为了平衡成本和用户体验，系统应该处理尽可能多的请求，直到系统负载达到预定义的阈值。在此之后，剩余的请求将被直接拒绝或推迟以进行后续重试。Mooncake 作为一个分解式推理系统实现，允许更灵活的调度策略，但也面临着非分解式系统中不存在的特殊调度挑战。
在系统过载的情况下，调度涉及根据系统负载确定是否接受或拒绝传入请求。此过程的一个关键方面是定义“系统负载”的构成，因为此定义会影响拒绝请求的阈值。在传统的耦合系统中，TTFT 和 TBT 的预测可能会因预填充和解码阶段之间的干扰而变得复杂。因此，负载通常仅通过正在处理的请求数与系统最大容量的比率来衡量。
相比之下，Mooncake 采用分解架构，独立处理预填充和解码阶段。因此，用 SLO 满意度作为直接负载测量。具体而言，分别将 lttf t 和 ltbt 定义为请求的 TTFT 和 TBT SLO 约束。然后通过将实例上的预测最大 TTFT 和 TBT 与 l/ttft 和 l/tbt 进行比较来确定预填充和解码实例的负载。基于这两个标准，Mooncake 的调度需要做出两个关键决策：第一，根据预填充实例的负载是否接受预填充阶段；第二，根据解码实例的负载是否继续进行解码阶段。
实际上，预填充或解码实例上的单个负载并不能准确反映系统处理的实际请求数。这种差异是由于单个请求的预填充和解码实例调度之间存在时间滞后而产生的。如果预填充阶段完成后，由于高负载而导致解码实例拒绝请求，则预填充阶段所消耗的计算资源将被浪费。因此，预填充期间成功处理的实际请求数小于负载指标所指示的数量。
为了解决这个问题，很自然地将解码实例的负载评估提前到预填充阶段开始之前。这种策略称为早期放弃（early rejection）。在请求到达时，Conductor 会根据预填充池和解码池之间的较大负载来评估是否接受请求。早期放弃显著减少了被拒绝请求造成的无效计算，并增强了负载平衡。
然而，早期放弃带来了新的挑战。如图显示了使用早期放弃策略后在 20 台机器 cluster 中观察的 20 分钟内实际实例负载。它突出显示预填充和解码机器之间的显著反相波动。这种现象在预填充机器较少的 cluster 和预填充阶段需要更长时间的场景中变得更加明显。
如图所示是应用早期放弃（a）和基于预测的早期放弃（b）的两种实例负载。绿色曲线表示预填充实例的负载（从 0 到 1 的尺度），黄色曲线表示解码实例的负载。
在第 1 阶段，预填充和解码实例上的负载都很低，因此 Conductor 会接受大量请求，直到预填充实例上的负载达到其极限。
在第 2 阶段，预填充实例处理的请求被调度到解码实例，导致解码实例的负载很高。因此，Conductor 拒绝传入请求，从而降低预填充实例的负载。
在第 3 阶段，没有新请求进入解码阶段，导致负载降低。此时，Conductor 再次接受大量请求，直到预填充实例满载。
在第 4 阶段，随着解码实例的负载增加，Conductor 拒绝请求，导致预填充实例的负载较低。预填充和解码实例之间的负载剧烈波动导致推理 cluster 的资源利用率低。
为了解决负载波动问题，提出一种基于预测的早期放弃框架，解决像 Mooncake 这样的分解式 LLM 服务系统在过载场景中的调度挑战。如上图（b） 所示，该框架预测传入请求预填充阶段后的解码负载，并使用此预测来决定是否接受请求，这有助于缓解波动问题。
该策略的核心部分是准确预测后续时期的解码负载。有两种方法：
请求级：以前的工作强调了预测 LLM 服务负载的重大挑战：每个请求的输出长度未知。如果可以提前确定输出长度，就可以更准确地估计 TTFT 和 TBT。这反过来又有助于预测解码实例可以完成的请求数以及在指定时间后将添加的新请求数，从而获得当时的负载。然而，由于成本高 [9] 或准确度低，预测每个请求的输出长度具有挑战性，尤其是在资源稀缺且需要准确预测的过载条件下，这使得请求级预测特别困难。
系统级：与请求级预测相比，系统级预测不会尝试预测单个请求的完成时间。相反会估计指定时间后实例的总体批次数量或 TBT 状态。这种类型的预测是持续进行的，对精度的要求较低，因此更适合过载场景。
在 Mooncake 中，目前采用系统级预测策略：假设每个请求的解码阶段都需要统一的时间 td。首先，对于给定时刻 t，将预填充实例在 t 时可以完成的请求添加到统一解码实例中。接下来，将在 t 之前完成（即其执行时间超过 td）的请求从解码实例中删除。最后，计算所有解码实例与 l/tbt 的平均 TBT 比率以预测负载。
版权归原作者 硅谷秋水 所有， 如有侵权，请联系我们删除。