月之暗面kimi底层推理系统方案揭秘(二)

极市平台

共 21301字,需浏览 43分钟

 ·

2024-07-11 22:00

↑ 点击蓝字 关注极市平台
作者丨ZHANG Mingxing
来源丨包包算法笔记
编辑丨极市平台

极市导读

 

深入探讨了月之暗面的底层推理平台—Mooncake。文章基于作者在知乎上的讨论和分享,特别是清华大学助理教授zhangmingxing的观点,详细介绍了Mooncake的分离式架构设计及其在实际应用中的优势和挑战。 >>加入极市CV技术交流群,走在计算机视觉的最前沿

既上一篇许欣然的月之暗面kimi底层推理系统方案揭秘,本篇继续。

作者分享在知乎上引起了广泛讨论,很多system方向的大佬炸出来了,本文根据作者清华助理教授zhangmingxing的一些分享整理,欢迎大家去围观知乎讨论区:

李博杰(华为天才少年):2020年我开始做分离式内存,做了非常高速的网络互联总线和很大的内存池(不能透露更多技术指标),当时搞了存储、数据库、AI和云混合部署等几个场景,感觉Persistent KV Cache会是disaggregated memory一个很重要的应用场景。很高兴看到Moonshot把它实现了。Prefill和Decoding分离,算力型硬件和内存带宽型硬件解耦,也是我非常喜欢的架构设计。

作者:昨天华为云的workshop上还聊到这个,分离式内存终于算是有了一个非常核心的落地场景。不过主要是带宽bound,iops bound 的场景想要落地还需要更多的探索,也包括 cxl/ub 之类硬件的进一步成熟

《Mooncake (1): 在月之暗面做月饼,Kimi 以 KVCache 为中心的分离式推理架构》

传送门:https://zhuanlan.zhihu.com/p/705754254

作者Disclaimer:和论文不同本文夹带过量私货个人观点,不代表本人单位更不代表 Moonshot,纯技术讨论。

TL;DR Mooncake 是由月之暗面创建的超人气智能助手 Kimi的底层推理平台。

本系列则是对应技术报告的浓缩版介绍。和强调 inclusion 四平八稳啥都讲一点的论文格式不同,这里更多的是想讨论一些当前还未形成共识的一些 design choice (私货警告)。特别的本篇主要讨论 Mooncake 的分离式架构,讨论点包括但不限于 TBT or TPOT,Prefill 节点应不应当独立存在以及独立的话如何多节点并行处理,KVCache cache 全局调度的原理和调度策略,Decode 还能不能进一步分离等等。

总体架构

Mooncake 的架构是非常典型的分离式架构,将单个同构 GPU 集群的资源打散并重新组织成三个可以独立弹性伸缩的资源池。其中 Prefill Pool 处理用户输入,主要对 Time To First Token (TTFT) 负责。同时因为 Prefill 相对计算密集,这一部分也承担着抬高整体资源利用率的任务。Prefill 处理完之后对应的 KVCache 会被送到 Decode Pool 进行 autoregression 式的流式输出。虽然我们希望尽可能攒大的 batch 以提升 MFU,但这一部分主要需要对 Time Between Tokens (TBT) 负责。

Mooncake 架构

这里我们使用 TBT 而非另外一些工作中常用的 Time Per Output Token (TPOT) 的原因在于虽然其实本来这两者应该等价,但实际计算 TPOT 的时候经常会被简单地等价成计算 average TPOT。然后就直接用了吞吐的倒数或者说总生成时间除以总生成 token 数。TBT 的定义则更加明确的指定为两个 token 生成间的延迟,因此基于 TBT 上限设计 Service Level Objective (SLO) 能更好的反映流式交互时的用户体验。直观的理解的话大家可以想象如果 GPT-4o 的语音回复是说一句停顿一段时间再说一句的说话方式的话,即便 average TPOT 能小一些体验还是不够好。

在 Prefill/Decode 之外我们还利用每台 HGX 机器上组成了一个 KVCache Pool 来进行全局的 Prefix Cache。相比 vLLM 当前单机的 Prefix Cache 通过全局的调度能够大幅度提升复用率从而提升总吞吐。由此带来了一系列如何调度,分配,复制 KVCache 的问题,而且必须和 Prefill/Decode 的调度 co-design,因此我们把 Mooncake 的架构称之为以 KVCache 为中心。

在此基础上接下来我们分别讨论 Prefill/KVCache/Decode Pool 的关键 design choice。

Prefill Pool

论文这一章我原本起了个“Prefill: To Separate or Not. Is it a Question?” 的名字。主要原因在于由于 Chunked Prefill 的提出实际上 Prefill 和 Decode 节点的边界已经模糊了起来。不过除非是特别短的 prompt (可以直接一次性加入到 decode 的 continues batch 里提升 MFU 但不违反 TBT SLO 的那种)不然我们还是倾向于使用独立的 Prefill 节点。原因主要是两方面。

首先我们定义了一个 VRAM occupation cost 的概念,即 KVCache 大小乘以其在显存中留驻的时间。使用 chunked prefill 的话由于每个 chunk 还需要和其他 decode request batch 在一起处理因此是显著增加了对应 KVCache 在显存中留驻的时间的,因此等效的降低了宝贵显存空间的利用率。当然显存空出来了有没有地方用就是另外一个问题了。这里我们在论文 4.2 节中主要讨论了和未来 batch job 结合的可能性。

另外更重要的一个原因则主要是在于长文本场景下我们需要在 Prefill 集群中使用特殊的多节点分布式划分方法来压低 TTFT。这里我们讨论了相比现在被讨论的更多的 SP 策略直接使用 Chunked PP 策略在调度和减少通讯量方面的优越性,前者大幅度简化架构,后者则能够省出宝贵的传输带宽来搬运 KVCache。具体策略其实就是 TeraPipe 的推理简化版本,因为只有 forward 没有 backword 的话不需要啥动态规划来平衡每个 Chunk 的计算量,只要给每个 Chunk 设置一个最小的值就行。

KVCache Pool

都 2024.6 了单节点的 Prefix Cache 已经不是什么新鲜东西了。这里主要讨论一下全局调度能够成立的底层原因:对于每 X byte 的 KVCache,其生成所需的算力正比于 X*hd 再乘以一个较大的常数,这里 hd 对应模型的 hidden dimension。因此只要每卡算力比如 A100 的 220TFLOPS 和每卡通讯带宽比如 CX7 的 100Gbps 的比值小于 hd 乘以这个常数,那么从远端传输 KVCache 相比原地重算不仅仅减少了计算量还减少了 TTFT。又省钱用户体验还更好了自然是何乐而不为。

但是上述公式中具体的 hd,常数是和模型结构强相关的所以不是说一定可以,只能说原则上比较大的模型应该都行。另外虽然网卡有 100Gbps,但是如果因为比如 KVCache 下沉到慢速存储上去了,网络拥挤了之类的原因跑不到这么高那么自然也就不一定成立。所以还是需要以 KVCache 为中心针对不同请求的 TTFT SLO 分别去调度。在论文的 5.2 节中我们主要是提供了一种简单的基于 heuristic 的热点识别和复制方法。倒不是说有多么精妙,但重点是简单好实现,不用特别去预测未来的 KVCache 使用模式,而且在实际的使用中效果不错的方案。未来会尝试更多的调度算法不过这个可以作为一个 strong baseline 来使用。

顺带这里感慨一句,目前标准的 8 卡 HGX 服务器每台会配 9~10 张 100Gbps 乃至 200Gbps 的 RDMA 网卡,至少 2TB 的内存和一大堆的高速 SSD。假设一千台的话就是近万张网卡互连起来的数 PB 共享内存池。做 Disaggregated Memory 挺久了第一次见这样的富裕仗。可惜作为一个 bandwidth bound 场景在架构上能设计的东西没那么多特殊的奇淫技巧可用,更多还是工程实现。

Decode Pool

Decode Pool 本身业界的实践是最多的。特别是开源有极快速地在迭代的 vLLM,该有的不是已经有了也是在 roadmap 里。这里主要是夹带一下 future work 中讨论的增加面向大容量大带宽设计的高性价比设备,形成异构集群的私货。从去年开始就断断续续和很多人讨论这个话题,虽然听闻了很多的风声,不过实际正用上的设备目测还得有个两年。而且硬件毕竟迭代慢,出来了算法变了(比如 MLA)或者被老黄的大力出奇迹创飞了就很尴尬。

不过实际上我们几个月前就挂出但没多少人看的另一篇论文也提到,其实不需要等新的硬件,只要组合基于 GDDR 的设备其实就能将 Decode 部分进一步拆分成 attention 算子和 Linear 算子两个池子从而进一步提升整体性能。不过目前云上的 GDDR 设备一般不会和旗舰显卡集群放在一起,而且为了节省成本只会配很差的 CPU 和网卡,所以想要实际部署这一套方案还有很多挑战。而且当前设备毕竟不是专门为这个场景设计的,attention offload 策略性价比的提升的潜力没有办法充分发挥出来,为此进一步将架构复杂化而且增大一定的 TBT 不一定很合适。所以归根结底还是得看看有没有硬件厂商愿意赌一把了。

其它

恭喜若愚今天正式以优秀毕业生的身份毕业,本科就这么厉害博士记得带我起飞~

《Mooncake (2):Kimi “泼天的流量”怎么接,分离架构下基于预测的调度策略 》

传送门:https://zhuanlan.zhihu.com/p/706204757

作者Disclaimer:和论文不同本文夹带过量私货个人观点,不代表本人单位更不代表 Moonshot,纯技术讨论。

TL;DR 正如 @许欣然 所说,本文公开的目的之一就是推动硬件厂商和云厂商向前篇提到的分离式,乃至未来异构分离的方向演化。作为我个人已知范围内承载了最大规模在线流量的分离式推理系统,Mooncake 也积攒了一些独特的实践经验。下篇这里主要是想要讨论其中非常关键的一项:Kimi 在承接急速增长的需求时,面向过载(Overload-oriented)场景的调度策略。

Kimi 也会累

和一般按照峰值容量规划,调度主要解决如何尽可能利用闲置资源的传统工作不同,Kimi 自从破圈之后面临的是完全不同的,如何在不断扩容还天天过载的情况下尽可能保障用户体验的问题。虽然是一个非常凡尔赛的甜蜜烦恼,同时 infra 也已经在最大程度避免,但在高峰期不可避免地还是需要偶尔向大家喊累。

这里主要是我们发现由于当前的负载情况已经无法按照 SLO 提供响应,所以与其提供一个完全无保障的长时间排队(违反 TTFT SLO)或者一字一顿的慢速回答(违反 TBT SLO)这样极大伤害体验的服务,不如先暂停一小部分服务。至于为什么是喊累而不是显示 429 Too Many Requests,当然是因为 Kimi 是大家的智能_伙伴_,伙伴累了要休息一下希望大家可以体谅。

分离式架构在过载场景下的挑战

也正是在这样的环境下正式上线了 Mooncake。虽然 Mooncake 大幅提升了 Kimi 能够承载的总吞吐,但 Prefill 和 Decode 分离的架构在过载调度上也的确引入了一些新的挑战。这里最重要的问题就是 Prefill 和 Decode 是分别对 TTFT 和 TBT 负责的,而且有一个关键的时间差。所以在坏的情况下 Prefill 前调度器觉得可以同时满足 TTFT 和 TBT,但是 Prefill 后由于 Decode 集群的过载 TBT 无法满足了。这样就进入了一个是违反 SLO 还是浪费已经花费了的 Prefill 资源的两难问题。

为了解决上述问题,一个自然的,也是被 Mooncake 采用的解决方案就是同时综合 Prefill 和 Decode 两个集群的情况,然后以两者中更高负载更危险的那个集群为去判定是否可以接受服务。由于会因为在 Prefill 集群有空闲的情况下由于未来 Decode 集群的负载问题提前拒绝一些服务,这样的策略被我们称之为 Early Reject。

看起来很直接,但实际部署之后我们观测到集群负载出现了奇怪的颠簸现象。可以看到 Prefill 和 Decode 集群的负载就和跷跷板一样一边上去一边下来,然后交替。

仔细分析之后发现的原因下图有一个很直观的展示。由于 Prefill 和 Decode 集群负载之间的时间差,如果简单的参考当前 Decode 集群负载去拒绝请求的话会导致 Decode 集群负载被消化的时候 Prefill 没有及时跟上,由此产生了跷跷板效应。

基于预测的调度

为了解决上述的问题,我们进一步设计和上线了基于预测的调度策略。原理也很直观,能够预测未来(特别是新加入的请求完成 Prefill 阶段的时刻)Decode 集群负载的话自然就可以打消时间差的问题平滑整个系统的负载。

具体来说,对于未来的某个时刻 t,首先我们将现有的 Prefill 集群中 t 时已完成的请求加入 Decode 集群,这一点可以通过预测 Prefill 时间和排队时间来实现。然后,我们假设每个请求的 decode 时间均为 t_d,将 Decode 集群中 t 时已结束(即decode 时间超过 t_d)的请求移除。这一步我们也在尝试通过预测每个请求实际会输出的 token 数做一个更精确的预测。最后,我们利用事先拟合的模型根据 Decode 集群中 t 时仍在处理的请求预测 TBT,作为 Decode 集群的预测负载并根据这个数据决定是否需要 Early Reject。

未来:调度会越来越复杂,也会越来越重要 从目前的趋势来看,未来 LLM Serving 的负载会愈发的复杂和多元化。原本同等地位的请求会在厂商和用户的驱动下基于多种多样的诉求进一步分化。比如有些批处理请求可以完全不考虑 TTFT/TBT 仅仅需要保障 Job Completion Time (JCT) SLO,有些高级/付费用户无疑需要保证更加可靠和更高的 SLO。另外外部基于 Agent 的任务也不仅仅看单个 request,而是看整个 Workflow 的综合 SLO,加上 Code Executor 等外部环境的交互后也会进一步复杂化整个系统。

在这些领域上我们也是刚刚起步,探索了一些,但是需要研究的更多,也希望借此机会和大家更多的探讨这一部分未来可能的优化空间。

技术报告原文原文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


公众号后台回复“数据集”获取100+深度学习各方向资源整理

极市干货

技术专栏:多模态大模型超详细解读专栏搞懂Tranformer系列ICCV2023论文解读极市直播
极视角动态欢迎高校师生申报极视角2023年教育部产学合作协同育人项目新视野+智慧脑,「无人机+AI」成为道路智能巡检好帮手!
技术综述:四万字详解Neural ODE:用神经网络去刻画非离散的状态变化transformer的细节到底是怎么样的?Transformer 连环18问!

点击阅读原文进入CV社区

收获更多技术干货

浏览 2
点赞
评论
收藏
分享

手机扫一扫分享

举报
评论
图片
表情
推荐
点赞
评论
收藏
分享

手机扫一扫分享

举报