极市平台 | 月之暗面kimi底层推理系统方案揭秘

本文来源公众号“极市平台”，仅用于学术分享，侵权删，干货满满。

原文链接：月之暗面kimi底层推理系统方案揭秘

极市导读

本文介绍了月之暗面科技有限公司开发的Kimi AI的底层推理系统方案Mooncake，以及在过载情况下的调度策略，包括早期拒绝和基于预测的早期拒绝策略，以提高资源利用率和系统稳定性。 

太长不看版（作者大佬自己的在知乎碎碎念）：本论文与很多 Prefill/Decoding 分离的论文不同的是，这套方案已经在大规模集群上进行几个月的验证并证明了方案的有效性。目前这套系统承载了 Kimi 线上80% 以上的流量， 效果很好也为产品带来了更多的设计空间。这也是为什么一个 POC 写在所有业内论文之前的系统，直到今天才发布出来跟大家见面。系统是需要跟随着应用快速变化的，同时也需要 硬件厂商 和 云厂商 早点接受新的理念才能跟上浪潮。

系统是需要跟随着应用快速变化的，同时也需要 硬件厂商 和 云厂商 早点接受新的理念才能跟上浪潮。发出这篇论文，主要是希望给各方提供一些信心，提供一些推理规模足够大场景下的必然优化思路。

趁这个机会，希望给各家硬件厂商和云厂商一些“暴论”

1.Mooncake 这类的存算分离策略会是一个长期趋势。

• 现在、立刻、马上真能省很多钱（毕竟不能公开规模和每日请求的 pattern，如果你说省不了那你都对）。

• KVCache 的容量会长期保持高位，因此围绕着 KVCache 来优化是非常有必要的。"private memory per request” 是整个推理系统优化中的关键瓶颈（否则 groq 万岁），会有很多努力来降低 KVCache 的大小，但同时会有更多动力来增大。

• 分离之后，允许整个系统往 “算力/”和“带宽” 的两个方向独立发展，对硬件优化是更友好的。AR 的模型架构惯性在短期内难以颠覆，因此总可以认定 decoding 的成本总会跟 bandwidth 成本有非常强的正相关性质。“带宽/低一个数量级的硬件方案已经在肉眼可见的范围内了。纵观历史，“算力/” & “带宽/$” 同时最优的芯片似乎还没出现过，集群上必然要拆分成异构的两个部分。

*2.Mooncake 方案和 MLA、各种 KVCache 的压缩方案 都是正交的，KVCache 变小了意味着 Mooncake 的方案收益会更明显。

*3.的确，有很多方向可能会让 Mooncake 的架构变成没必要的方案。

• given hardware lottery，新架构演进会是一个相当缓慢的过程（不 AR、用 RAG 做 Attention 等等方案），不能因噎废食。

• 包括我们自己也在投入很多资源在 break 现有框架上，因此有理由相信在可见的未来推理方案还会变动。

• 由于目前海量的推理压力，所以软件系统做为一个迭代速度极快的方案，就应该一代模型一次跟进。

• 预测这个状态至少会持续2~3年，因此集群层面现在已经值得做拆分了。

• 芯片层面值得做为一个重要的设计考量，在芯片的 IO 能力上要多预留一些能力。

Github地址：https://github.com/kvcache-ai/Mooncake/tree/main

原文（powered by kimi/互相赋能就让它自己宣传自己吧）

1 引言

1.1 动机：开发 Mooncake 的原因

随着大型语言模型（LLMs）在各种场景中的迅速采用，对LLM服务的工作负载已经变得显著多样化。这些工作负载在输入/输出长度、到达频率和分布上有所不同，最重要的是，它们需要不同种类的服务级别目标（SLOs）。作为一个模型即服务（MaaS）提供商，Kimi [5] 的一个主要目标是解决一个具有多个复杂约束的优化问题。优化目标是在满足不同级别的SLOs的约束下，最大化整体有效吞吐量，这直接影响收入。

为了实现这一目标，一个先决条件是充分利用GPU集群中可用的各种资源。具体来说，尽管当前GPU服务器以高度集成的节点形式提供（例如，DGX/HGX超级计算机 [6]），但有必要将它们解耦并重新结构化为几个分离的资源池，每个池针对不同但协作的目标进行优化。例如，许多研究人员 [7, 8, 9] 已经建议将预填充服务器(prefill servers)与解码服务器(decoding servers)分开，因为LLM服务的这两个阶段具有非常不同的计算特性，在请求从预填充服务器转移到解码服务器时，KVCache会发生变化。

基于这个想法，我们发现KVCache的调度是LLM服务调度的核心。为了提高整体吞吐量，通常有两种一般方法：1）尽可能多地重用KVCache以减少所需的计算资源；2）最大化每个批次中的token数量以提高模型FLOPs利用率（MFU）。然而，从远程位置重用KVCache会延长首次token的时间（TTFT），而较大的批次大小会导致更大的token间时间（TBT）。因此，这两种面向吞吐量的优化的使用可能会导致违反与延迟相关的SLOs。

根据上述指导方针，我们提出了一个以KVCache为中心的分离设计，用于调度和优化。图1展示了我们当前的以KVCache为中心的分离架构，名为Mooncake。对于每个请求，全局调度器（Conductor）需要选择一对预填充和解码实例，并按以下步骤调度请求：1）尽可能多地将可重用的KVCache转移到选定的预填充实例；2）以块/层的方式完成预填充阶段，并将输出的KVCache连续地流式传输到相应的解码实例；3）在解码实例中加载KVCache，并将请求添加到连续批处理过程中以生成请求输出。

尽管这个过程看起来简单，但由于许多限制，选择策略却相当复杂。在预填充阶段，主要目标是尽可能多地重用KVCache以避免冗余计算。然而，等待存储在低级存储上的KVCache可能会导致违反TTFT SLO。此外，对KVCache服务器的高需求可能会导致网络拥塞，延长等待时间。因此，Conductor还负责预测KVCache块的未来使用情况，并相应地执行调度操作，例如交换和复制。最热门的块应该复制到多个节点以避免获取拥塞，而最冷的块应该被交换出去以减少保留成本。预填充调度还受到预填充节点上DRAM空间可用性的限制，特别是当大量内存被保留用于全局KVCache池时。

相比之下，解码阶段有不同的优化目标和约束条件。目标是将尽可能多的token聚合到一个解码批次中以提高MFU。然而，这一目标不仅受到TBT SLO的限制，还受到可以包含在VRAM中的聚合KVCache总大小的限制。

更重要的是，现有的LLM服务研究假设资源充足，并专注于提高资源利用率。相比之下，当前的GPU/加速器供应有限，许多MaaS提供商面临严重的过载问题，尤其是在高峰时段。在这种情况下的调度提出了现有工作尚未探索的独特挑战。例如，我们需要预测未来的负载，并在预填充阶段后如果没有可用的解码插槽，则尽早拒绝某些请求，以节省浪费的计算资源。然而，这种早期拒绝策略的直接实现令人惊讶地导致了过载的波动。这促使我们致力于预测特定查询的生成长度，并在短期内进行整体负载预测，以实施更好的拒绝策略。同样必要的是，对不同请求优先级进行分类，以实施基于优先级的调度。在本文中，我们将这些问题总结为面向过载的调度，并展示了我们的初步研究结果。

1.2 Mooncake 的设计与结果

在本文的后续部分，我们首先介绍 Mooncake 架构的概览，包括其主要组件和处理请求的典型工作流程（§3）。然后，我们描述了在实现过程中做出的主要设计选择，特别是那些当前研究中未涵盖的部分。

首先，在§4中，我们讨论了如何实现一个单独的预填充节点池，以无缝处理上下文长度的动态分配。我们采用了分块流水线并行机制（Chunked Pipeline Parallelism, CPP）来扩展单个请求在多个节点上的处理，这对于减少长上下文输入的首次Token时间（TTFT）是必要的。与传统的基于序列并行（Sequence Parallelism, SP）解决方案相比，CPP减少了网络消耗，并简化了对频繁弹性扩展的依赖。这种机制通过层级预填充进一步补充，使KVCache的流式传输能够重叠延迟。

接下来，在§5中，我们详细阐述了我们的以KVCache为中心的请求调度算法，它平衡了实例负载和通过首次Token时间（TTFT）和Token间时间（TBT）SLOs测量的用户体验。这包括一个基于启发式的自动热点迁移方案，它在不需要精确预测未来KVCache使用情况的情况下复制热门KVCache块。实验结果表明，我们的缓存感知调度可以显著降低现实世界场景中的TTFT。在使用公共数据集、模拟数据和真实工作负载的端到端实验中，Mooncake 在长上下文场景中表现出色。与基线方法相比，Mooncake 能够在满足SLOs的同时实现高达525%的吞吐量增加。在真实工作负载下，Mooncake 使 Kimi 能够处理75% 的更多请求。

最后，与假设所有请求都将被处理的现有LLM服务工作不同，由于Kimi 用户请求的快速增长，Mooncake 持续面临过载问题。因此，Mooncake 的调度涉及根据系统负载决定是否接受或拒绝传入请求。在§6中，我们讨论了我们独特的早期拒绝策略的实现，它减少了过载场景中浪费的计算资源。我们进一步探讨了由简单早期拒绝引起的负载波动问题，以及如何通过预测未来负载来缓解这个问题。

Mooncake 目前是服务 Kimi 的主要平台，并已成功处理了呈指数级增长的工作负载，证明了其在扩展到大型和高度过载工作负载方面的有效性。然而，还有更多问题需要探索，这些未来的发展方向也包括在本文中。为了保护专有信息并促进可复制性，本文报告的所有实验结果都是基于真实工作负载的重放跟踪，但使用的是一个遵循与LLaMA2-70B相同架构的虚拟模型。该跟踪仅包括请求到达的时间、输入token的数量和输出token的数量，不包含任何真实用户内容。在遵循某些内部程序后，该跟踪将稍后开源。

2 预备知识和问题定义

现代大型语言模型（LLMs）基于Transformer架构，该架构利用注意力机制和多层感知器（MLP）来处理输入。基于流行的Transformer模型，例如GPT [10] 和 LLaMA [11]，采用的是仅解码器结构。每个推理请求在逻辑上被划分为两个阶段：预填充阶段和解码阶段。

在预填充阶段，所有输入token并行处理。此阶段生成第一个输出token，同时存储计算出的中间结果，这些中间结果被称为KVCache。解码阶段随后使用这个KVCache来自回归地生成新的token，将新计算出的键和值添加到KVCache中。预填充阶段能够同时处理输入token，通常使其在计算上非常密集，除了短请求之外。由于注意力网络的计算复杂度与输入长度呈二次方增长，如图2左侧所示，预填充阶段的计算时间通常随着输入长度的增加而超线性增长。

与此相反，解码阶段由于自回归生成的限制，每个批次只处理一个token。这使得它受内存限制，并且计算时间随着批次大小的增加而次线性增长，如图2右侧所示。在解码阶段广泛使用的优化是连续批处理 [12, 13]。在每次迭代之前，调度器检查所有请求的状态，将新到达的请求添加到批次的预填充阶段，同时移除已完成的请求。

由于预填充阶段和解码阶段的不同特性，MaaS提供商设置了不同的度量标准来衡量它们各自的服务级别目标（SLOs）。具体来说，预填充阶段主要关注请求到达和生成第一个token之间的延迟，即首次Token时间（TTFT）。另一方面，解码阶段侧重于同一请求连续token生成之间的延迟，即Token间时间（TBT）。

作为一个MaaS提供商，通过满足服务协议中定义的SLO度量标准来确保质量保证至关重要。例如，TTFTP 90 = 4×这样的度量标准表明，在相同条件下无干扰地运行的单个请求的TTFT的90%不超过四倍。具体来说，在本文的端到端实验（§7.1）中，我们设置了TTFTP 90 = 10×和TBTP 90 = 5×。在真实部署中，我们设置了固定的TTFT和TBT SLOs。如果监控检测到未满足的SLOs，我们要么增加推理资源，要么拒绝一些传入的请求。

然而，由于当前GPU的有限供应，通常不可行地扩展推理集群。因此，在面向过载的调度中，决定拒绝哪些请求成为核心问题。我们的主要目标是在遵守SLOs的同时最大化整体吞吐量，这一概念在其他研究中被称为goodput [8, 14]。我们的方法不同之处在于，只有完全完成执行的请求才会计入goodput的度量。否则，所有之前消耗/生成的token都不会被计算，相应的资源就会被浪费。换句话说，如果一个请求在SLO下无法完成其全部执行，则应尽早拒绝该请求。

实现这一目标不仅需要优化预填充和解码阶段的架构，还需要开发预测短期未来负载的能力。

3 Mooncake 分离架构的概述

正如图1所示，Mooncake 采用了一种分离架构，它不仅将预填充(prefill)节点与解码(decoding)节点分开，而且还将GPU集群的CPU、DRAM、SSD和RDMA资源分组，实现了一个分离的KVCache。这种分离缓存利用了未充分利用的资源，提供了充足的缓存容量和传输带宽，使得在不增加额外成本的情况下，能够高效地实现接近GPU的前缀缓存。

KVCache池在CPU内存中的存储和传输逻辑如图3所示。在CPU内存中，KVCache被存储为分页块。根据请求模式，它可以采用如最近最少使用(LRU)、最少频繁使用(LFU)等缓存逐出算法，或基于请求特征的算法。这些KVCache块在CPU和GPU之间的传输由一个称为Messenger的独立(GPUDirect) RDMA基础组件处理。此架构还使我们能够为外部用户提供上下文缓存API，以实现KVCache的更高重用。

为了调度所有这些分离组件，在Mooncake的中心，实现了一个名为Conductor的全局调度器。Conductor负责根据当前的KVCache分布和工作负载来调度请求。如果对未来推理有利，它还会复制或交换某些KVCache块。具体来说，图4展示了一个请求的典型工作流程。一旦完成令牌化，Conductor就选择一对预填充节点和一个解码节点，并启动一个包含四个步骤的工作流程：

1) KVCache重用： 被选定的预填充节点（组）接收一个请求，该请求包括原始输入、可以重用的前缀缓存块ID，以及分配给请求的完整缓存块ID。它根据前缀缓存块ID从远程CPU内存加载前缀缓存到GPU内存以启动请求。如果没有前缀缓存存在，则跳过此步骤。此选择平衡了三个目标：尽可能重用KVCache、平衡不同预填充节点的工作负载，并保证TTFT SLO。这导致了将进一步讨论的以KVCache为中心的调度。

2) 增量预填充： 预填充节点（组）使用前缀缓存完成预填充阶段，并将新生成的增量KVCache存储回CPU内存。如果未缓存的输入token数量超过某个阈值(prefill_chunk)，预填充阶段将被分成多个块，并以流水线方式执行。此阈值被选择以充分利用相应GPU的计算能力，通常大于1000个token。使用分块但仍然分离的预填充节点的原因在§4.1中解释。

3) KVCache传输： 前述的Messenger服务在每个节点中部署，用于管理和传输这些缓存。每个Messenger在其各自的推理实例中作为一个独立进程运行，接收信号以促进高速跨机器KVCache传输。这一步是异步执行的，并与上述增量预填充步骤重叠，将每个模型层生成的KVCache流式传输到目标解码节点的CPU内存中，以减少等待时间。

4) 解码： 在解码节点的CPU DRAM中接收到所有KVCache后，请求以连续批处理的方式加入下一个批次。Conductor根据其当前负载预先选择解码节点，以确保它不会违反TBT SLO。然而，这种SLO会由本地调度器再次检查，因为预计的负载可能在预填充阶段之后发生了变化。这种双重检查可能导致请求被拒绝，在这种情况下，相应的预填充成本就会被浪费。

4 预填充池的实现

与不可违反的解码节点不同，设计一个单独且弹性的预填充池的必要性和最佳实践仍然存在争议。例如，尽管许多研究人员[7, 8, 9]与我们有相同的直觉，即使用分离架构，但在引入分块预填充[15]后，分离是否仍然必要仍然值得讨论。分块预填充将输入token分成多个小块，这些小块加入连续批处理过程。这种方法有两个明显的好处：1) 无需分离，所有节点被视为平等的，使调度变得更容易；2) 将分块预填充内联到解码批处理中可以提高解码批处理的计算强度，从而实现更好的模型FLOPs利用率(MFU)。

然而，在仔细考虑后，我们决定保持Mooncake的分离架构。只有在不需要分块且不影响TBT SLO的情况下，请求的预填充才会内联到解码批处理中。做出此决定有两个主要原因：1) 预填充节点需要不同的跨节点并行设置来处理长上下文(§4.1)；2) 它提供了节省VRAM的独特机会(§4.2)。

4.1 多节点预填充

最近LLMs的可用上下文长度正在迅速增加，从8k到128K甚至1M[16]。通常，对于这样的长上下文请求，输入token可能是输出token的10到100倍，这使得优化TTFT变得至关重要。由于长上下文预填充中存在大量的并行性，因此使用超过一个8x GPU节点来并行处理是可取的。然而，将张量并行性(TP)扩展到超过一个节点需要每层进行两次昂贵的基于RDMA的all-reduce操作，显著降低了预填充节点的MFU。

最近，许多工作提出了序列并行性(SP)[17, 18, 19, 20, 21, 22, 23]。SP通过在不同节点上分配请求的输入序列来实现加速。这些SP方法利用了注意力运算符的关联属性，并且至少需要在每层实现Ring Attention[18]或Striped Attention[19]期间进行一次跨节点通信。这大大减少了网络消耗并提高了MFU。

然而，即使采用SP，与仅使用单节点TP相比，MFU仍然更差。理想的部署是将预填充节点组织成两组：一组仅使用TP，另一组使用SP。仅在必要时为满足TTFT SLO时才将请求调度到SP组。这种进一步的分离导致了动态调整每个组中节点数量的问题，因为静态并行设置可能导致集群的低利用率。最近的研究[14]提出了弹性序列并行性来动态扩展或缩减SP组。尽管这是可能的，但它为我们的架构增加了复杂性。例如，它需要提前建立一个全局通信组，并在考虑如缓存重用利用率和SLO要求违规等指标时，使Conductor的设计复杂化。这使得我们的系统在需要频繁即时可伸缩性的情况下面临挑战。此外，SP仍然需要频繁的跨节点通信，这降低了MFU，并与跨节点传输KVCache的网络资源竞争。

为了解决这个问题，Mooncake利用了仅解码器变换器的自回归属性，并为长上下文预填充实现了分块流水线并行性(Chunked Pipeline Parallelism, CPP)。我们将预填充集群中的每X个节点分组为一个流水线预填充节点组。对于每个请求，其输入token被划分为块，每个块的长度不超过prefill_chunk。相同请求的不同块可以由不同的节点同时处理，从而并行化处理并减少TTFT。

CPP提供了两个主要好处：1) 与训练中的流水线并行性类似，它仅在每个流水线阶段的边界处需要跨节点通信，这可以容易地与计算重叠。这带来了更好的MFU和更少的网络资源争夺与KVCache传输。2) 它自然适应短上下文和长上下文，为短上下文预填充带来了没有显著开销的好处，并避免了频繁的动态节点划分调整。这种基于流水线的加速方法已经在训练系统中被探索[24]，但据我们所知，这是在推理阶段的首次应用，因为长上下文推理直到最近才出现。

4.2 逐层预填充

除了计算能力外，VRAM的有限大小也是宝贵的资源，我们的目标是通过状态，主要是KVCache，最小化VRAM的占用。理论上，如果一个请求的KVCache大小为S，处理时间为T，则其占用成本为S*T。如果一个请求被分块，并且每个块的处理与解码请求在分块预填充中内联，则T将增加，导致更大的占用成本。

此外，由于预填充是逐层处理并且受计算限制，因此可以将KVCache的传输和转储与计算重叠，进一步降低其占用成本。在Mooncake中，KVCache的加载和存储通过启动和等待操作异步执行。在每层的注意力计算开始之前，模型等待该层的KVCache异步加载完成，并触发下一层的异步KVCache加载。在注意力计算完成后，启动该层KVCache的异步存储。一旦所有层的计算完成，进程等待所有异步存储操作的完成。传输重叠允许预填充实例的执行时间大致等同于KVCache加载时间或标准预填充时间，这取决于相对于输入长度的前缀缓存比例。

实验结果，如图5所示，KVCache存储延迟表明，逐层预填充可以有效地减少长上下文请求的延迟。这种重叠效果的主要优点是，它使我们能够在预填充调度中忽略可用VRAM的大小，只要它能够容纳一个请求。如图1所示，预填充节点的调度只考虑KVCache的分布和可用的DRAM大小。

将来，我们打算探索这种空闲VRAM的更多用途。例如，OpenAI最近提出了使用批量API[25]，它使用户能够以50%的低成本发送异步请求组，但只有明确的24小时周转时间。此服务非常适合处理不需要立即响应的工作。由于这些批量请求没有严格的TBT，我们可以在有足够的VRAM空间容纳相应的KVCache的情况下，甚至将这些请求的解码阶段内联到预填充处理中以获得更好的MFU。

5 以KVCache为中心的调度

在本节中，我们主要讨论在正常条件下，Conductor如何调度请求和KVCache块，下一部分将讨论过载场景下的讨论。

5.1 预填充全局调度

以往关于LLM服务的研究通常使用基于每个实例分配请求数量的负载均衡策略。然而，在Mooncake中，选择预填充实例时会考虑额外的因素——不仅是负载，还有前缀缓存命中长度和可重用KVCache块的分布。虽然倾向于将请求路由到具有较长前缀缓存长度的预填充实例以减少计算成本，但有时为了确保整体系统平衡和满足TTFT SLO，将它们调度到其他节点可能是有益的。为了解决这些复杂性，我们提出了一个考虑前缀缓存导致的预填充时间和实例负载相关的排队时间的缓存感知全局调度算法。

算法1详细说明了我们的缓存感知预填充调度机制。对于每个新请求，其输入token被分成几个块，并为每个块计算一个哈希键。这涉及到为一个块中的token生成一个哈希键，并将该哈希键与前一个块的哈希键（如果有）连接起来。然后，将请求的块键与每个预填充实例的缓存键逐一比较，以确定前缀匹配长度（prefix_len）。类似的重用逻辑已经在vLLM中实现，但vLLM的开源版本仅支持本地KVCache缓存。

有了这些匹配信息，Conductor根据请求长度和prefix_len（因实例而异）估计相应的执行时间。然后，它将该请求的估计等待时间加到该实例的TTFT上。最后，Conductor将请求分配给具有最短TTFT的实例，并相应地更新该实例的缓存和队列时间。如果无法达到SLO，Conductor直接向上层返回HTTP 429 Too Many Requests响应状态码。

这个调度框架的核心是直接的，但是各种组件的工程实现中隐藏着复杂性。例如，为了预测请求的预填充阶段的计算时间，我们采用了一个基于离线测试数据的预测模型。该模型根据请求的长度和前缀缓存命中长度估计预填充持续时间。由于Transformer的计算模式是规则的，只要有足够的离线数据，这个预测的误差范围就很小。一个请求的排队时间是通过聚合所有排队请求的预填充时间来计算的。在实际实现中，TTFTs是并行计算的，使得处理时间与推理时间相比可以忽略不计。

预测传输时间更加困难，因为它不仅由传输数据的大小决定，还由当前的网络状态决定，特别是发送节点是否处于拥塞状态。这也需要复制热门KVCache块，这将在下一节中讨论。

5.2 缓存负载均衡

在我们的Mooncake集群中，每台预填充机器管理自己的一组本地前缀缓存。这些缓存的使用频率差异很大。例如，系统提示几乎被每个请求访问，而存储本地长文档内容的缓存可能只被一个用户使用。如§5.1所讨论的，Conductor在实现缓存匹配和实例负载之间的最佳平衡中起着至关重要的作用。因此，从分布式缓存系统的角度看，负载均衡也起着重要作用。具体来说，它涉及策略制定，如何备份缓存以确保全局预填充调度能够实现高缓存命中率和低负载。

解决这个KVCache调度问题的一个初步解决方案可能是收集每个块的全局使用情况，使用预测模型预测它们未来的使用情况，并据此做出调度决策。然而，与预填充时间的估计不同，工作负载是高度动态的，并且随时间显著变化。特别是对于用户基数快速增长的MaaS提供商来说，准确预测未来的使用情况是不可能的。因此，我们提出了一个基于启发式的自动热点迁移方案来增强缓存负载均衡。

正如前所述，由于高实例负载，请求可能不会总是被定向到具有最长前缀缓存长度的预填充实例。在这种情况下，Conductor将缓存的位置和请求转发到一个替代实例，如果估计的额外预填充时间短于传输时间。这个实例主动从持有者检索KVCache并将其存储在本地。更重要的是，我们更倾向于计算输入token，如果最佳的远程前缀匹配长度不大于当前本地可重用前缀乘以一个阈值。这两种策略不仅减少了请求的预填充时间，还促进了热点缓存的自动复制，允许它们在多台机器上更广泛地分布。

6 面向过载的调度

大多数现有的关于LLM服务的研究都假设所有请求都将被处理，相应地优化吞吐量或请求的TTFT和TBT。然而，在实际情况下，处理每一个传入的请求既不经济也不现实。对于面临用户请求量迅速增长的商业推理服务来说，集群的推理资源增长速度远远落后于传入请求的增加。因此，过载是当前LLM服务中常见的问题，尤其是在高峰时段。

为了平衡成本和用户体验，系统应该尽可能多地处理请求，直到系统负载达到一个预定义的阈值。达到这一点后，剩余的请求将被直接拒绝或推迟到以后再重试。作为实现为分离式推理系统的Mooncake，允许更灵活的调度策略，但也面临非分离式系统所没有的独特调度挑战，且在以前的工作中没有提及。

在这一部分，我们描述了一个为分离式架构特别设计的早期拒绝策略，并解决了这种方法引起的负载波动问题。然后，我们探讨了预测生成长度对于缓解这些问题的必要性。

6.1 过载场景中的调度

在系统过载发生的情况下，调度涉及到根据系统负载决定是否接受或拒绝传入的请求。这一过程中一个关键方面是定义什么构成“系统负载”，因为这一定义影响请求被拒绝的阈值。在传统的耦合系统中，由于预填充和解码阶段之间的干扰，预测TTFT和TBT可能变得复杂。因此，负载通常简单地通过正在处理的请求数量与系统最大容量的比率来衡量。

与此相反，Mooncake以其分离式架构处理预填充和解码阶段的独立性。因此，我们使用SLO满足度作为直接的负载测量。具体来说，我们定义了lttft和ltbt分别为请求的TTFT和TBT SLO约束。然后，预填充和解码实例的负载通过将实例上预测的最大TTFT和TBT与lttft和ltbt进行比较来确定。有了这两个标准，Mooncake的调度需要做出两个关键决策：第一，根据预填充实例的负载是否接受预填充阶段；第二，根据解码实例的负载决定是否继续解码阶段。

6.2 早期拒绝

在实践中，预填充或解码实例的个别负载并不准确地反映系统处理的实际请求数量。这种差异源于调度单个请求的预填充和解码实例之间存在时间滞后。如果一个请求在预填充阶段完成后由于解码实例的高负载而被拒绝，那么在预填充阶段花费的计算资源就被浪费了。因此，预填充阶段成功处理的请求实际数量少于负载指标所指示的数量。

为解决这个问题，很自然地将解码实例的负载评估提前到预填充阶段开始之前。我们将此策略称为早期拒绝。在请求到达时，Conductor根据预填充和解码池之间的较大负载评估是否接受请求。早期拒绝显著减少了被拒绝请求的无效计算，并提高了负载均衡。

6.3 早期拒绝引起的负载波动

然而，早期拒绝引入了新的挑战。图7显示了在使用基于预测的早期拒绝策略后，在20分钟内观察到的20台机器集群的预填充和解码实例负载。它突出了预填充和解码机器之间的显著反相波动。当预填充机器较少，以及预填充阶段较长的场景中，这种现象更加明显。

经进一步探索，我们发现这种负载波动问题根源于预测解码负载与实际执行之间的时间差。基于当前解码负载的调度本质上存在延迟。这个延迟导致预填充和解码实例的负载出现波动和相位错位，如图8a中描述的理论示例所示。绿色曲线代表预填充实例的负载（按0到1的比例缩放），黄色曲线代表解码实例的负载。

在第一阶段，预填充和解码实例的负载都很低，因此Conductor接受了大量的请求，直到预填充实例的负载达到其限制。在第二阶段，由预填充实例处理的请求被调度到解码实例，导致解码实例的负载很高。因此，Conductor拒绝传入的请求，导致预填充实例的负载降低。在第三阶段，没有新的请求进入解码阶段，导致负载减少。此时，Conductor再次接受大量的请求，直到预填充实例完全负载。在第四阶段，随着解码实例的负载增加，Conductor拒绝请求，导致预填充实例的负载降低。这种预填充和解码实例之间负载的严重波动导致推理集群的资源利用率低下。

6.4 基于预测的早期拒绝

为解决负载波动问题，我们提出了一个面向过载场景的分离式LLM服务系统（如Mooncake）的基于预测的早期拒绝框架。如图8b所示，该框架预测了预填充阶段后传入请求的解码负载，并使用此预测来决定是否接受请求，这有助于缓解波动问题。该策略的核心组件是准确预测后续期间的解码负载。我们为此引入了两种方法：

请求级别：以前的工作强调了预测LLM服务负载的一个重大挑战：每个请求的未知输出长度。如果我们能够提前确定输出长度，就有可能更准确地估计TTFT和TBT。这反过来有助于预测解码实例可以完成的请求数量，以及在指定时间后将添加的新请求数量，从而获得该时间点的负载。然而，由于成本高昂[9]或准确度低，尤其是在资源稀缺和需要准确预测的过载条件下，使得请求级别的预测特别困难。

系统级别：与请求级别的预测不同，系统级别的预测不尝试预测单个请求的完成时间。相反，它们估计在指定时间后实例的整体批次计数或TBT状态。这种类型的预测是持续的，并且需要的精度较低，使其更适合过载场景。

在Mooncake中，我们目前使用系统级预测策略：我们假设每个请求的解码阶段需要统一的时间td。首先，对于给定时刻t，可以由t时刻的预填充实例完成的请求被添加到统一的解码实例中。接下来，在t之前可以完成（即，执行时间超过td）的请求从解码实例中移除。最后，计算所有解码实例的平均TBT比率与ltbt的比值来预测负载。将请求级预测的探索留作未来的工作。

7 评估

本节评估了Mooncake在不同数据集和各种工作负载下的端到端性能。如前所述，为了保护专有信息并促进可复制性，本文报告的所有实验结果都是基于一个虚拟模型，该模型遵循与LLaMA2-70B相同的架构。

7.1 端到端性能

表1: 在端到端实验中使用的数据集。

数据集	平均输入长度	平均输出长度	缓存比率	到达模式
ArXiv摘要[26]	8088	229	~0%	泊松过程
L-Eval[27]	19019	72	>80%	泊松过程
模拟数据	16k, 32k, 64k, 128k	512	50%	泊松过程
真实数据	7955	194	~50%	基于时间戳

7.1.1 公共数据集

本节评估了Mooncake和vLLM在公共数据集上的端到端测试性能，使用了ArXiv摘要和L-Eval。我们使用四个vLLM实例作为基线，表示为vLLM-[4M]。相比之下，Mooncake有两种不同的配置：一个集群由三个预填充实例和一个解码实例组成，标记为Mooncake-[3P+1D]；另一个有两个预填充和两个解码实例，标记为Mooncake-[2P+2D]。图9显示，在ArXiv摘要和L-Eval数据集上，Mooncake-[3P+1D]在满足SLO的同时，分别比vLLM-[4M]实现了20%和40%的吞吐量提升。此外，L-Eval数据集上Mooncake的吞吐量通过前缀缓存进一步提高，显著减少了预填充时间。然而，尽管具有较低的TBT延迟，Mooncake-[2P+2D]在TTFT指标上的表现不如Mooncake-[3P+1D]和vLLM-[4M]。这种差异源于预填充和解码实例之间的负载不平衡。在现实世界的集群中，预填充和解码实例的需求通常在一定时期内保持稳定，只有轻微的暂时性不平衡。因此，预填充和解码实例的比例可以预设。未来的研究将探索更灵活的部署和转换方法。

7.1.2 模拟数据

本节使用模拟数据进行端到端实验。集群配置与§7.1.1相同，使用Mooncake配置[3P+1D]、[2P+2D]和vLLM-[4M]。值得注意的是，模拟数据中的长上下文请求显著破坏了vLLM的解码阶段。为了对抗这一点，vLLM单独处理请求，而不是批量处理。实验结果如图10所示。尽管Mooncake采用批量处理，但其两阶段分离设计有效地最小化了预填充阶段对解码阶段的影响，确保它从未违反TBT SLO。Mooncake展示了显著更高的吞吐量，提升范围从50%到525%，同时遵守与vLLM相同的TTFT和TBT SLO约束。

7.1.3 真实工作负载

我们进一步利用10个预填充实例和10个解码实例，标记为Mooncake-[10P+10D]，以及20个vLLM实例，称为vLLM-[20M]，重放真实请求跟踪并在Mooncake和vLLM上进行负载测试。在这个实验设置中，TTFT的上限设置为30秒，而TBT阈值限制在每个token 0.1秒。图11展示了两个系统在TTFT和TBT上的累积分布函数(CDF)图。Mooncake-[10P+10D]和vLLM-[20M]的TTFT分布几乎相同，几乎所有请求都满足了TTFT SLO。然而，尽管Mooncake-[10P+10D]的所有请求都满足了TBT SLO，但vLLM-[20M]的请求中只有57%满足了这一标准，一些请求显示出极高的TBT。在这个实验中，Mooncake可以在遵守SLO的同时处理大约75%的更多请求。

7.2 过载场景下的性能

本节评估了在过载场景下的性能，重点关注系统可以处理的最大请求数量，如§6所述。基线策略在两个阶段开始之前基于负载拒绝请求，导致资源浪费，因为已经处理过的预填充阶段的请求被拒绝。与此相反，我们提出了早期拒绝和基于预测的早期拒绝策略，分别在§6.2和§6.4中详细说明。这些策略综合考虑了系统的负载，从而减少了不必要的请求拒绝。

具体来说，我们构建了一个具有8个预填充实例和8个解码实例的Mooncake集群，并使用23,000个请求的真实跟踪对其进行了测试。为了模拟过载场景，我们将重放速度提高到2倍。

表2显示了Mooncake在不同策略下的系统拒绝请求的数量。使用基线策略时，系统拒绝了4183个请求。相比之下，在早期拒绝和基于预测的早期拒绝策略下，Mooncake分别拒绝了3771和3589个请求。这表明，通过早期拒绝请求，Mooncake可以避免不必要的预填充计算，从而提高系统资源的有效利用率。此外，通过预测解码实例的负载，Mooncake可以缓解负载波动，增加请求处理能力。

策略类型	拒绝请求的数量
基线(Baseline)	4183
早期拒绝(Early Rejection)	3771
基于预测的早期拒绝(Early Rejection based on Prediction)	3589

THE END !

文章结束，感谢阅读。您的点赞，收藏，评论是我继续更新的动力。大家有推荐的公众号可以评论区留言，共同学习，一起进步。