|
随着大型语言模型的社会影响力日益增强，相应的人工智能产品用户基数也在迅速扩大。目前，AI 产品发展的一个主要挑战是如何在有限的计算资源下，有效应对日益增长的用户需求。 在 2024 年 10 月 18−19 日举办的 QCon 全球软件开发大会（上海站） 上，月之暗面推理系统负责人何蔚然分享了“Mooncake 分离式推理架构创新与实践”，他从实际业务出发，讲述了在固定集群资源的条件下，通过采用单点和分布式推理架构，提升集群处理大规模请求的能力的挑战和解决思路，希望能给大家带来一些帮助。 内容亮点：
以下是演讲实录（经 InfoQ 进行不改变原意的编辑整理）。 大家好，我是月之暗面的推理系统负责人何蔚然。今天，我非常荣幸能与大家分享我们在分离式推理架构领域的一些探索和所面临的挑战。我们构建的这个框架被命名为 Mooncake。接下来，我将从四个维度展开介绍。 首先，我们将探讨大规模推理中的挑战，尤其是在处理长上下文场景时，我们需要解决哪些问题。其次，我们会讨论单点性能优化，即在单个实例上，我们如何进行性能提升。第三，我们将深入探讨分离式架构，分析在单实例优化的基础上，如何通过分离式架构实现更高效的性能提升。在此过程中，我们也会分享一些关键的观察，包括为什么现有的一些方案可能并不适用。最后，我将展望未来，在硬件和软件方面探讨未来的发展趋势，并期待与更多的合作伙伴共同探索这一领域。 我们公司的拳头产品名为 Kimi 智能助手，以及 API 形式提供服务的 Kimi 开放平台。在 Kimi 智能助手的体系中，还有一系列专注于特定领域的应用，我们称之为 Kimi plus。这些应用包括塔罗师、长文生成器等，它们服务于不同的业务场景，各自的负载需求也不尽相同。尽管如此，所有这些不同的业务都依赖于同一套推理引擎来提供支持。每天这些子业务系统会产生数以万亿计的 token 量，这要求我们的处理能力必须极为强大。我们的负载特点主要是长上下文处理，这是我们产品定位的一个结果。在这种背景下，我们对用户有着严格的服务水平目标（SLO）保证。为了满足这些保证，我们的集群长期处于满载状态。因此，我们面临的挑战是如何在不过载的情况下优雅地处理更多的用户请求。为了应对这一挑战，我们采取了一些特别的并行和调度策略。通过这些策略的实施，我们的成本相比去年已经大幅下降超了 20 倍。
首先我分享一下公司对于推理成本降低的核心价值观念。我们坚信，随着时间的推移，推理成本必然会降低，但我们模型的智能水平不能因降本而下降。如果为了降低成本而牺牲模型的智能水平，用户不会满意，我们的工作也将失去意义。基于这个前提，我们发现可选择的路径并不多，主要有两个方向：一是提高算子的计算速度，二是降低算子或算法对显存的需求，从而实现更高的并行处理能力。或者，我们可以寻找性价比更高的硬件，以替代目前价格高昂的显卡。 我们总结了几个关键公式：
这里所说的更节省的模型结构，指的是在时间和显存上的优化。如果我们进一步拆分推理成本，会发现两个关键点：一是长上下文的 Prefill（预填充），二是 Generation（生成）的成本。对于长上下文的预填充，我们知道 Attention（注意力机制）具有平方级的时间复杂度。随着上下文长度的增加，比如达到 64K 或一兆，所需的时间也会呈平方级增长，这成为我们系统中非常关键的一部分。因此，我们需要对这种场景进行专门的优化。优化长上下文预填充后，我们发现从整体上看，生成的成本才是推理系统的主要成本。因为用户在对话过程中需要模型输出的字数越来越长，而生成是一个 Memory Bound 的过程。 对于长上下文，要使其成本更低，要么提高注意力机制的计算速度，要么减少 KV 缓存的大小，从而提高并行度。对于生成，我们需要更大的批量大小来均摊成本。由于其内存密集型的特性，算子执行时间相对较短，因此我们需要避免通信成为生成阶段的瓶颈，需要对 Decode（解码）过程采取更友好的并行方式。 长上下文处理面临的挑战主要有两个方面。首先是 Full Attention（全注意力机制）的时间复杂度问题，它是 n 的平方，这意味着对于像 Llama 3 这样的经典模型结构，进行 1 兆的注意力计算可能需要几十分钟。对于多轮对话，如果 像开源框架选择在每一轮对话中都重新计算完整的上下文，将浪费大量的 GPU 时间，因为这些计算是可以被重复利用的。其次，随着上下文长度的增加，KV Catch 会占用更多的显存，这限制了我们同时处理的请求量，无法达到很高的并行度。如果我们不采用 Prefill 和 Decode 分离的架构，就需要为 Prefill 的峰值显存预留空间，这会影响 Decode 的 Batch Size（批量大小）。这两点最终都会导致无法有效扩大 Batch Size，从而无法实现成本的有效均摊。 除了性能挑战，我们还需要在大规模推理时采用一些自动运维手段，以减少人力投入，专注于解决更重要的问题。为此，我们采取了以下措施：首先，我们实现了推理实例的快速切换和快速拉齐方法。由于显卡是容易损坏的硬件，我们有硬件巡检手段，能够在机器出现问题时快速隔离，并在一定时间内如果无法恢复则人工介入。其次，在深夜时段，推理压力不大时，我们会释放一部分空闲资源来执行一些长时间或离线的任务，这些任务对延迟不敏感，可以异步进行。或者将这些机器用于一些轻量级的训练任务，以避免资源闲置。 在单点性能优化方面，无论是 Prefill 还是 Decode，我们都实施了一系列措施。首先，我们采用了混合并行策略，旨在为 Decode 设计更优的并行方式。下面是我们几种策略：
前三种并行策略在当前的开源框架中都得到了较好的支持。特别是 PP，通常用于显存不足或需要处理离线任务的场景。然而，在在线场景中，我们有更优的选择 CP，即 Context Parallelism（Ring/AllToAll），它能有效分摊长文本的算力需求。我们从序列维度进行切分，有两种实现方式：Ring CP 和 AllToAll CP。Ring CP 允许计算和通信较好的重叠，从而节省通信成本。但如果我们使用的是 Sparse Attention（稀疏注意力）机制，每个 Rank 之间的计算强度不等，这时 Ring CP 可能无法实现完美的重叠，因此我们推荐使用 AllToAll CP 实现。 AllToAll CP 虽然能解决 Sparse Attention 的问题，但它引入了新的挑战，即通信时间无法重叠。为了解决这个问题，我们采用了 Chunked Pipeline Parallelism（CPP），它能够在超长文本推理中完美掩盖计算和通信。我们可以想象，如果有两个负载相同的请求，我们可以在一个请求进行通信时执行另一个请求的计算，这样显卡在整个时间片上都能保持高负载，没有资源浪费。但这种方法会导致每个请求的延迟变长，对系统的压力增大。因此，我们采用了论文中提到的另一种方式，将一个请求巧妙地拆分成许多计算强度接近的部分，利用每个 Chunk 在通信时进行下一个 Chunk 的计算，实现计算和通信的完美重叠。 CP 和 CPP 在 Prefill 阶段非常有用，但在 Decode 阶段，由于序列长度通常不会太长，这两种方法不太适用，甚至可能引入额外的通信成本。因此也采用了 Data Parallelism（not LB），即单卡完整计算单组请求以减少通信代价。在一般的框架中，DP 可能被理解为请求之间的负载均衡，但我们这里指的是减少 Decode 的通信，使单组请求能在单张卡上完整计算，从而节省时间延迟。 最后，我们的多种并行方式可以共享同一个通信组。例如，在训练工作中，TP8 加 CP8 可能意味着需要 8 台机器进行计算，但在推理时实际上不需要。因为 TP 适合 MLP 时使用，而 Attention 适合用 CP8。这时，我们可以选择在同一个 8 卡组中，前半部分使用 CP，后半部分使用 TP。 在长上下文推理优化方面，我们开展了一系列工作：
首先，我们有 Moonshot Sparse Attention 技术，这是我们经常被人们问及的无损长文压缩技术。今天可以明确告诉大家，我们不使用 Full Attention，但我们的技术仍然能够实现无损压缩。其工作原理是在 Attention Mask 上不进行完整的计算，从而减少计算量并降低延迟。此外，这种技术还能减少 KV Cache 的大小，使我们能够提高并行度。 接下来是 Dynamic MoE 技术。我们发现在每个请求的所有输出 token 中，所谓的 Easy tokens 占比较高。对于这些 Easy tokens，我们不需要激活那么多的 Expert 就能完成高效的推理。因此，由于计算量的减少，我们能够提高 batch size。 在 Speculative Decoding 方面，许多研究者都在探索如何共享 KV Cache IO。这是我们各种投机推理工作的主要收益点。最近，我们还发现类似的工作可以用于共享 expert IO，未来我们将以论文形式分享这一发现。 Cascade Attention 是我们在产品中常见的一种功能，例如在回复一段话后自动生成标题或推荐问题等。这些功能背后的原理是将上下文与特定的 System Prompt 结合，让模型针对特定任务输出结果。这些任务的共同特点是需要处理较长的上下文，但输出长度不会太长，因为标题可能只有几个字。对于这种场景，如果任务量增加，上文的 KV Cache IO 开销会相当大。因此，我们可以选择将这些任务集中处理，只需支付一次 KV Cache IO 的代价，就能覆盖所有这些短输出。 分离式架构是本次分享我重点要讨论的内容。我们为什么要采用这种架构？在实施过程中遇到了哪些挑战？首先，我们可以观察到 Prefill 和 Decode 是两种截然不同的负载场景。Prefill 是一个明显的计算密集型负载，而 Decode 则是一个明显的内存密集型场景。对于 Prefill，我们追求的是最大化 MFU ，即每人民币算力性价比最高的硬件。对于 Decode，我们需要的则是每人民币带宽性价比最高的硬件。我们还发现，去年市场上的显卡定价在这两类负载上存在错位。利用这个错位，我们可以进一步降低成本。 我们的设计场景要求严格保证 SLO，并且我们更倾向于处理长上下文。在这两个条件下，我们需要最大化吞吐量。我们的 SLO 主要有两个方向：第一个是 TTFT（Time to First Token），即用户输入一段话后，需要等待多长时间才能收到第一个字的反馈。如果这个时间过长，用户可能就会失去耐心。第二个是 TBT（Time Between Tokens），即 token 生成的速度必须稳定地快于人类的阅读速度。 针对这些观察，我们最终发现背后的关键是如何尽可能节约计算资源。这最终归结为一个技术—— Prefill Cache ，即我们能共享的上文越多越好，这样计算资源就能节省得越多。我们总结了几种有效场景：
我们的优化工作带来了显著的收益，这些收益体现在几个关键指标上。首先，我们实现了 TTFT 的 10 倍提升，这主要得益于 Cache Miss 的显著降低，目前我们能做到小于 10% 的 Cache Miss 水平，大量的计算可以被重复利用。 其次，我们在 TBT 上获得了大约 5 倍的提升。这主要归功于 decode 节点能够将 batch size 增大两倍以上。如果我们采用 Prefill 和 Decode 混合部署的方式，Decode 节点的 TVT（Time to Value）压力会比较大，因为需要在 Decode 之间插入 Prefill 的计算。但如果我们将 Prefill 和 Decode 分离， Decode 节点就不需要为 prefill 预留任何显存，从而可以增大 batch size。尽管如此，batch size 的增加也会导致 TBT 相应下降，因此在 SLO 的限制下，我们最终只能达到两倍多的水平。 在总体吞吐量上 RPM 上，我们平均获得了 1.7 倍的提升，对于一些较简单的业务，提升甚至超过了 5 倍。这些成绩的取得，是因为我们挖掘了当前许多框架可能没有充分利用的硬件资源，例如基础架构的 RDMA 通信带宽、内存的容量和带宽，以及 OSS 或 SSD 等多级缓存工具。 我们的架构设计包含三个主要部分：Prefill 节点池、Decode 节点池和一个虚拟的 KV Cache 池。架构图的左上角蓝色部分代表 Prefill 节点池，右下角橙色部分是 Decode 节点池。中间绿色的部分是 KV Cache 池，它通过 RDMA 连接 Prefill 和 Decode 节点，而其管理则由左侧灰色区域的 KV Cache 中心调度器负责，它决定每个请求应该在哪个 Prefill 机器上执行，以及在哪个 Decode 机器上进行生成，同时调度 KV Cache 的传输时机。
因为 Prefill 和 Decode 的负载特点不同，所以我们可以在 Prefill 节点池选择算力较强的显卡，而在 Decode 节点池选择带宽较好的显卡，实现异构推理显卡的部署。由于负载特点的差异，我们可以为 Prefill 和 Decode 选择不同的并行策略，这意味着 Prefill 中的 KV Cache 布局与 Decode 中的布局不同，这种差异需要由中间的 KV Cache 池来弥补，使得 Decode 节点无需了解 Prefill 的 KV Cache 具体结构。 我们还利用了 RDMA 的带宽和内存的容量带宽。有人可能会质疑，内存容量和带宽已经有框架做得很好了，比如今年各个推理框架都增加了一个功能叫做 Local Prefix Cache。这种方法在一些基准测试中非常有效，甚至可能超越我们的系统。但在业务量极大时，本地内存很容易被耗尽，一旦耗尽，Cache Miss 会大幅上升，性能急剧下降。因此，在实际的海量请求业务系统中，这种方法并不实用。幸运的是，我们有 RDMA 硬件，可以利用它来传输内存中的 KV Cache，它能突破 Local Prefix Cache 的限制，达到更好的 Cache Miss 效果。 关于 SSD 和 OSS 这类存储方式，虽然它们的带宽可能不高，甚至较慢，许多公司选择不使用它们存储 KV Cache。但实际上，只要它们不违反 SLO，我们就可以利用它们的带宽来节省 GPU 处理时间。 许多公司在尝试实现 P&D 分离后发现效果并不理想，我们分析可能存在以下两个原因。 首先，Prefill 的 Cache Miss 问题可能导致恶性循环。我们的请求通常是一波一波到来的，有时我们需要同时处理多个热门请求。以单机为例，任何机器都有一个并发上限。当超过这个上限时，我们必须选择其他机器进行请求分流。这时，如果机器 B 接到了一个热门请求，但机器 B 上没有相应的 KV Cache，那么它唯一的选择就是重新计算，这将消耗大量时间。反过来，这会增加机器 B 的 TTFT 压力，限制了全局并发的并发度。因此，一波请求可能导致整个系统的 Cache Miss 迅速上升，形成一个恶性循环：需要解析的越多，可用并行度越低，而并行度越低，又需要解析的越多，导致性能急剧下降。 其次，Decode 对延迟非常敏感。因此，我们需要在机内和机间进行计算和通信的重叠，以使通信延迟尽可能不可见。我们希望在 Decode 机器上，显卡在任何时刻都在进行计算，这是我们最理想的状态。 我们的两种方法最终都指向了一个核心问题：如何有效地传输 KV Cache？为了解决这个问题，我们分别采取了以下措施。 首先，针对 Prefill Cache Miss 的问题，关键在于机器 B 没有热请求的 KV Cache。我们的解决方案是采用 Prefill 到 Prefill 的 Cache Transfer。当机器 B 发现没有 KV Cache 时，我们不选择重新计算，因为这会消耗大量时间，而是让机器 A 直接将 KV Cache 传输给机器 B。这样，机器 B 就可以打破恶性循环，减少 TTFT 压力，提高并行度。 其次，我们需要处理 Prefill 到 Decode 的传输，这使得我们的 RDMA 网络带宽使用非常频繁。因此，我们需要一个更优的 RDMA 传输方案。许多开源工具的实现可以达到 80GB 的水平，但离理论上限还有一定距离。我们对 RDMA 传输进行了精细调整，使得传输速率可以达到 180GB 每秒，非常接近 200GB 的理论上限。 RDMA 通信结束后，我们面临的是一个 RAM 到 RAM 的过程。在实际计算时，我们需要将 KV Cache 从内存搬运到显存，以及将计算完的 KV Cache 从显存搬运回内存。这个过程很容易成为瓶颈。我们有两种方法来解决这个问题。第一种是利用 Copy Engine，通过共享内存进行中转，但这种方式的带宽利用率只有 30% 左右。我们采用的是另一种方法，根据计算强度动态调整显卡的 Streaming Multiprocessor，直接从内存搬运到显存，这样带宽利用率可以提高到 80% 以上，接近理论上限。 即使有了这些技术，仍然不够。想象一下，如果我们面对的是如 Llama 3 这样的 65B 级别的模型，它有 80 层的 KV Cache 需要传输。一般的做法可能是 Prefill 计算完成后，再决定开始向 Decode 机器传输，这会在瞬间产生巨大的传输压力。为了避免这种情况，我们采取了一种方法，就是每一层的 KV Cache 计算完成后立即发出。因为从整个时间线来看，这类事件是非常稀疏的，只是由于一次性传输太多而变得显著。 我们的讨论转向 KV Cache 中心的调度器设计，这是我们架构中进行调度优化的核心部分。所有这些优化本质上都是为了更好地利用重复计算。以下是我们采取的一些方法。 首先，我们实现了 Prefill&Decode 节点之间的动态平衡。这是因为线上负载并非恒定不变，有时 Prefill 的压力会更大，有时则是 Decode。如果采用固定配比，硬件利用率不会是最优的。我们有两种实现方式。如果 Prefill 节点池和 Decode 节点池共享同一硬件，我们可以直接进行低延迟的状态切换，同时保留内存上的 KV Cache，避免重传。如果 Prefill 节点和 Decode 节点是异构硬件，我们则需要通过水平扩缩容来实现动态平衡。 其次，我们处理了许多相同的多轮对话或衍生的多轮对话。如果仅靠负载均衡器进行分发，Cache Miss 会很严重。因此，我们尽量让这些请求集中在同一台机器上计算，利用 Cascade Attention 共享 IO。 我们实现了热点均衡，即让热的 KV Cache 尽量多地分配到不同机器上，而无需人工运维，保持系统健康。我们采用了启发式方法，根据节点的当前负载和 KV Cache 命中长度计算得分，选择得分最高的机器进行调度。这是因为我们可以预测请求的输出长，从而在请求到达后，就能在调度层面安排其整个生命周期，包括在哪个系统上做 Prefill，哪个系统上做 Decode，以及何时传输 KV cache。 这种提前调度的方法使得调度器的 CPU 开销得以隐藏，因为在 Prefill 过程中，GPU 计算时间较长。此外，有了输出长度的预测后，我们可以做出更好的拒绝策略。如果系统已满载，对于无法处理的新请求，我们会立即拒绝，以免浪费计算资源。 我们有一些开源计划。首先是 Trace Dataset，这是一个我们开源的数据集（ https://github.com/kvcache-ai/Mooncake）。它的意义在于，我们之前提到的许多性能数字都可以通过这个数据集得到复现。我们将线上负载进行脱敏处理后公开，这样社区成员可以利用这些数据实现自己的调度算法或策略，看看是否能够实现比我们更高的命中率。 其次是 Mooncake Store，这是一个统一的 KV Cache 文件系统，目前还未开源，敬请期待。KV Cache Transfer 是一个特别重要的技术手段，其效率至关重要，且需要进行复杂的并行或异步调度。这个过程中会涉及到许多复杂的细节，比如不同级别的存储需求不同，不同运营策略的存储格式也不同。Mooncake Store 就是为了帮助大家解决这些问题，使开发者能够更专注于如何实现更优的调度策略。同时，我们希望这个项目能够成为 KV Cache 池化接口层的一个标准。如果云服务提供商希望与我们合作，我们非常欢迎。
展望未来硬件的发展，我认为有几个关键点。首先，由于生成仍然是成本的主要部分，因此内存和显存的容量与带宽越大越好，或者性价比越高越好，这将有助于我们进一步降低成本。目前，我们还没有充分利用 MBU（Memory Bandwidth Utilization），如果我们能够将 MBU 提升至接近 90% 的水平，成本将显著下降。 其次，我们关注的是 KV Cache 的内存池化存储系统。我们希望未来的硬件能够允许我们跳过 RAM 到 RAM 的过程，直接从远端内存搬运到本地显存。 在分析过程中，我们发现 Prefill 和 Decode 的负载复杂性不同，进一步思考后，我们意识到 Attention、MLP 和 MoE 的计算强度也不同。因此，我们可以进一步分离这些组件，为 Attention 选择显存性价比更高的硬件，而为 MLP 或 MoE 选择计算性价比更高的硬件。我们相信这种精细分离能够实现降低成本的目的。 何蔚然，月之暗面推理系统负责人。毕业于清华大学计算机系，8 年以上云 + 端性能优化经验。目前在月之暗面负责自研推理框架开发，致力于智能平权，让更多用户体验到最高的智能水平。曾在旷视科技公司主力参与自研芯片、低位宽神经网络研发。  后续我将通过微信视频号，以视频的形式持续更新技术话题、未来发展趋势、创业经验、商业踩坑教训等精彩内容，和大家一同成长，开启知识交流之旅 |
|
|
