演讲嘉宾｜唐飞虎
编辑 ｜蔡芳芳
策划 ｜AICon 全球人工智能开发与应用大会
在不久前举办的 AICon 全球人工智能开发与应用大会上，月之暗面高级研发工程师、开发者关系负责人唐飞虎发表了专题演讲“长文本大模型推理实践——以 KVCache 为中心的分离式推理架构”，分享介绍 Kimi 智能助手背后的推理加速方案，以及该方案在设计时所需要考虑的指标和在真实生产环境中部署的表现。
在 10 月 18 -19 日即将召开的 QCon 上海站上，我们专门策划了【大模型基础设施与算力优化】专场，并邀请到月之暗面推理系统负责人何蔚然进一步分享 Mooncake 分离式推理架构创新与实践，同时微软亚洲研究院软件开发工程师姜慧强将分享 《长文本 LLMs 推理优化：动态稀疏性算法的应用实践》，还有更多大模型训练推理的一手实践案例尽在本专题。欲了解更多精彩内容，可访问大会官网：https://qcon.infoq.cn/2024/shanghai/schedule
以下是演讲实录（经 InfoQ 进行不改变原意的编辑整理）。
提到 Kimi，相信在座的各位都有所耳闻。Kimi 智能助手在多个平台上都有入口，包括 Apple Store、微信小程序以及 Web 端，尤其是 Web 端的排名一直居高不下。在日常使用中，尤其是在午间高峰时段，用户可能会遇到 Kimi“累了”的情况。但值得注意的是，尽管用户数量在不断增加，用户体验却得到了显著改善，现在 Kimi“累了”的情况减少了很多，这与我们推理团队的技术攻关是分不开的。
今天，我将从四个方面进行介绍。第一部分我将探讨长文本推理的瓶颈问题。随着推理集群的扩大和上下文长度的增加，无论是训练还是推理都面临着更高的要求。我们需要明确瓶颈所在，以便找到解决问题的途径。第二部分将审视目前市面上的推理优化工具和方法，看看有哪些可以为我们所用。第三部分将详细介绍我们的 Mooncake 项目。这是一个以 KVCache（键值缓存）为中心的分离式推理架构，内部代号“Mooncake”由我命名，寓意与“moon”相关，同时“Cake”与“Cache”谐音。在最近几个月，我们也看到了不少类似的方案提出，虽然与我们的设计大同小异，但我将分享一些细节，特别是在面对大量用户时可能遇到的一些独特问题。第四部分，我将讨论上下文缓存的应用。在 SaaS 服务层面，我们为开发者提供了上下文缓存功能。我将具体介绍每位开发者如何利用 Mooncake 方案来优化自己的 AI 应用。
长文本推理的瓶颈
几个月前，我在 AICon 北京站上讨论了 RAG 与长文本处理的对比。RAG 模型有其优势和劣势，但在长文本处理方面，它有两个显著的劣势：成本高和速度慢。这也是许多开源模型无法良好支持或只能有损支持长文本处理的原因之一。
成本高的问题是我们的用户群体，特别是 API 开发者在生产环境中经常遇到的一个痛点。他们需要使用我们的 API 对同一个文档，比如一份合同，进行复杂的任务处理，并可能需要反复询问 100 多次或更多。众所周知，大型模型通常采用无状态设计，这意味着每次调用都需要将整个上下文传递进去。随着对话的进行，上下文可能会不断增长，不仅包括上下文本身，还可能包括函数调用、定义以及其他文档设置。每次调用都需要处理这么多信息，因此成本自然会很高。
速度慢问题，特别是在第一次处理时，模型的响应速度会特别慢。例如，我们群里的 API 助手在用户频繁提问时，有时会卡住，20～30 秒内都无法产生回复。我们无法确定是卡住了还是其他地方出了问题，但很可能仅仅是因为处理速度慢而导致的卡顿。这是长文本处理速度慢带来的一个副作用。
贵且慢的原因
长文本推理之所以成本高昂且速度缓慢，原因在于 Transformer 模型在计算 Attention 机制时的工作方式。在没有使用缓存的情况下，每次计算 Attention 都需要进行完整的矩阵乘法，这导致每次迭代的长度以平方级别增加。每当出现一个新的 Query Token，都需要重新计算，这无疑增加了计算的复杂性和时间。
当我们引入 KVCache 机制后，情况就大为改观。使用 KVCache 后，每次计算的长度只需要线性增加，这意味着我们不再需要重新计算过往的 tokens，从而显著提升了性能。但这种优化也带来了新的问题。
为了更直观地理解这一点，我们可以看下面这张图。图表的横坐标表示上下文长度，从左到右逐渐增大；纵坐标则展示了不同的性能度量指标，包括并发数、预填充延迟、解码延迟以及上下文切换时的状态切换延迟。此外，还有 Free HBM Size，它表示可用的高带宽内存大小，随着上下文长度的增加，它可以改善并发性能。
在 Mooncake 论文中，我们使用了 LLaMA2-70B 模型作为例子。从图表中可以明显看到，随着序列长度（即上下文长度）的线性增长，预填充延迟（prefill latency）呈现出超线性增长，这是一个非常显著的趋势。
我们做个小结，长文本性能瓶颈主要包括以下几个方面：
并发性能：随着上下文长度的增加，并发性能会反比下降；
预填充延迟：随着上下文长度的增长，预填充延迟会以平方级别的速度增长；
解码延迟和上下文切换延迟：随着上下文长度的增加，解码延迟也会线性增加。
长文本推理的优化
面对长文本推理中的成本高昂和速度缓慢的问题，我们采取了一系列优化策略。在之前的讨论中，我们已经提到了一些方法，现在我们来总结一下这些策略，并探讨它们如何帮助我们解决长文本推理的挑战。
首先，我们考虑了几种优化技术，包括 Flash Attention、vLLM（垂直扩展的大型语言模型）、MOE（Mixture of Experts，专家混合模型）以及最近非常受欢迎的 Speculative Decoding（推测性解码）。此外，还有一些有损策略，例如 Windows Attention，它通过截取部分信息来处理长文本推理，尽管这种方法在资源有限的情况下可行，但最终提供给用户的模型可能是有损的。例如，如果用户询问一个公司上市财报的问题，信息可能分散在文档的不同部分，使用 Windows Attention 策略可能会遗漏一些关键信息。这种信息的遗漏是在训练过程中就已经决定的，后续的 SFT 也无法修复这个问题。虽然这种策略可能在短期内提高模型的上下文指标，但对最终用户体验来说可能是有损害的，而且这种损害在进行大模型基准测试时可能不容易被检测出来。
此外，这些策略之间存在不兼容问题。这就需要我们深入了解每个策略的具体实现方式，这会带来一系列复杂的问题，需要从不同的层、头或隐藏层角度进行推理优化。在系统设计时，我们需要选择兼容性最好且效果最佳的策略。
我们的 Mooncake 主要是从集群调度的角度进行优化。这种优化与我们之前提到的所有策略基本上是正交的，可以与任何策略组合使用，而不会损失性能。这对于模型基础供应商来说可能是一个高优先级的策略。这也是为什么我们最近看到许多厂商推出了基于 KVCache 优化的方案。这些方案能够提高长文本推理的效率，同时保持用户体验的高质量。
Mooncake 的实践
上个月，我们在 GitHub 上发布了相关的论文，其中包含了许多细节。在详细介绍之前，让我们先澄清一些基本概念。在大模型推理中，有两个至关重要的阶段：预填充（Prefill）阶段和解码（Decode）阶段。
预填充阶段：在这个阶段，每个新 Token 的生成都依赖于之前所有的 Token。由于输入的全部内容（即 Prompt）都是已知的，这个阶段可以进行高度并行化的矩阵操作，有效提高 GPU 的利用率。这个阶段对“首次 Token 时间”（Time to First Token, TTFT）有显著影响，对于流式应用来说，如 GPT-4o 或 AI 陪伴类实时交互应用，这个指标尤为重要。
解码阶段：与预填充阶段不同，解码阶段的 Token 不是一次性全部生成的，而是逐段、逐词生成的。这个过程会一直持续到满足某个停止条件，这些条件可能是模型的最大上下文限制、用户设置的上下文上限，或者是其他预设的停止条件，比如输出了一个 JSON 的终止符。在这个阶段，每个顺序输出的 Token 都需要知道之前所有迭代的输出状态的 KV 对。这涉及到矩阵中的向量运算，与预填充阶段相比，解码阶段无法充分利用 GPU 的计算能力。数据从内存传输到 GPU 的速度决定了这个阶段的延迟，而不是计算本身的速度。换句话说，解码阶段主要受内存传输速度的限制，这个阶段主要影响的是“每个输出 Token 的时间”（Time per Output Token, TPOT），对于流式应用来说，这个指标同样敏感，并且对总体推理时间有较大影响。
Mooncake 的基本思想
Mooncake 的核心理念是将模型推理过程中的两个截然不同的优化阶段分开处理，因为这两个阶段的优化目标和受限的瓶颈各不相同。这种分离式处理方法是一种直观且自然的思路。具体来说，Mooncake 采用了以 KVCache 为中心的分离式推理架构，主要由三个核心部分组成：
Prefill 池：这个部分负责集中管理所有的预填充阶段的计算任务。
Decoding 池：这个部分集中处理所有解码阶段的任务。
KVCache 池：这个部分负责存储所有中间过程中应用到的 KVCache，并决定何时使用这些缓存，何时释放它们。
Mooncake 的开发动机
现有的大型语言模型（LLM）服务系统，例如垂直扩展的大型语言模型（vLLM），通常将预填充（Prefill）和解码（Decode）阶段放在同一个 GPU 上处理。这种设计在实际操作中会遇到一些问题：
TTFT 和 TPOT 优化不可兼得：预填充阶段处理的文档通常非常长，并且随着多轮对话的进行，上下文长度会不断增加，导致预填充比解码耗时更长。在需要作出决策先执行哪个阶段时，vLLM 的做法是让解码阶段暂停，以便预填充阶段先行。这会导致 TPOT（每个输出 Token 的时间）增加。反之，如果优先解码，TTFT（首次 Token 时间）也会增加。
GPU 资源竞争：即使预填充和解码阶段被单独调度，它们仍然会竞争 GPU 资源，导致等待时间增加。在单机情况下，很难设计出一个高效的调度策略来解决这种资源竞争问题。
Prefill 和 Decode 的瓶颈不同：预填充阶段主要消耗计算资源，而解码阶段则更依赖于内存和带宽。目前还没有既擅长计算又擅长内存和带宽的芯片，因此，使用具有不同特点的 GPU 分别处理预填充和解码，可能更有利于资源的利用和成本的节约。
这种分离设计的思路在计算机体系结构领域并不新鲜。例如，区块链系统 EOS 在设计时就采用了类似的分离策略，将带宽和 CPU 计算分开处理，各自形成一个资源池。受到这种设计思路的启发，Mooncake 采用了典型的分离式架构，将单个同构 GPU 集群的资源打散并重新组织成三个可以独立弹性伸缩的资源池 —— Prefill Pool、Decode Pool、KVCache Pool。
Mooncake 的分离式架构
在 Mooncake 架构中，Prefill Pool 中的一个具体实例负责处理特定的任务。下图的工作流程可以这样理解：
KVCache 的重用（Reuse）：在下图左侧黄色区域，我们可以看到处理过程中可能会涉及到 KVCache 的重用。这部分 KVCache 可能已经存在，它包含了之前迭代中计算得到的数据，可以在新的迭代中被再次利用，从而提高效率。
增量 KVCache（Incremental KVCache）：在下图左侧粉色部分，是处理新的输入所生成的增量 KVCache。用户在每一轮对话中都可能提供新的输入，这些输入需要被处理并生成新的 KVCache，作为当前轮次迭代的结果。
流量式传输（Traffic-style Transfer）：计算完成后，这些 KVCache 会以一种流量式的方式传输，即连续不断地传递给解码阶段的实例。
解码实例（Decoding Instance）的任务：接收到 Prefill Pool 传递过来的 KVCache 后，解码实例的任务就相对简单了。它主要负责根据这些 KVCache 进行解码操作，生成最终的输出结果。
在 Mooncake 架构中，KVCache 资源池扮演着至关重要的角色，尤其是在线上服务中，每天需要应对数百万用户的查询。在一些热点事件，如奥运会期间，用户可能会频繁询问类似的问题，比如中国队当天获得的金牌数量或者乒乓球男团的晋级情况。面对这样的高频查询，我们可以利用 KVCache 资源池来优化处理。
一个自然的想法是使用哈希存储来管理这些 KVCache。不同的厂商可能会采用不同的策略，例如有些可能会选择使用 Trie 树。Trie 树在计算复杂度上与模型的词表大小有关，如果模型的词表发生变化，比如从 GPT-4 升级到 GPT-4o 时词表扩大了，Trie 树的性能可能会受到影响。
为了避免这种复杂性并提高效率，我们选择了哈希存储。哈希存储方法简单、速度快，并且不受词表大小变化的影响。这样，无论用户的查询如何变化，我们都能保证 KVCache 资源池的高效运作。
Mooncake 效果展示
在实施了三个独立的资源池之后，我们进一步观察了线上生产环境中的实际运行情况。我们发现，Prefill 和 Decoding 阶段对资源的占用呈现出一种波浪型模式，类似于潮汐的涨落，每天都有规律地变化，尽管看起来似乎没有明显的规律。然而，当我们仔细观察并分析每天线上数百万用户的行为时，这些行为模式变得可以预测。基于这些可预测的行为模式，我们可以采取类似于 vLLM 中的策略，即在资源紧张时暂停解码阶段，让预填充阶段先行。在集群调度中，我们也可以应用类似的逻辑。我们根据线上生产环境的日常运行数据，设计了一种基于动态规划的调度策略。这种策略能够提前准备适量的预填充和解码资源，以使资源占用的波动曲线更加平缓。通过这种动态规划的调度策略，我们能够更有效地管理资源，减少资源浪费，并确保服务的稳定性和响应速度。这样的调度提升了用户体验，因为我们能够更好地应对用户需求的高峰和低谷，确保服务始终如一地流畅运行。
在下面展示的图表中，蓝色线条代表我们的 Mooncake 架构，而黄色线条则代表与之对比的 vLLM 架构。我们可以观察到两种架构在首次 Token 时间和 Token 间时间上的表现差异。
首次 Token 时间：在首次 Token 时间上，Mooncake 架构相较于 vLLM 架构有轻微的优化。这主要是因为引入了 KVCache，它能够存储和重用之前的计算结果，从而加快了预填充阶段的处理速度。
Token 间时间：在 Token 间时间上，Mooncake 架构的改善非常明显。这是因为 Mooncake 将解码阶段单独分配给一个资源池来处理，这样的分工使得解码过程的效率显著提高。由于解码阶段不再与其他任务共享资源，它可以更加专注和高效地生成每个后续的 Token。
下图展示了一些具体的实验数据，这些数据并非来自线上环境，而是使用了开源数据集。我们特别关注了 ArXiv 上每天发布的论文摘要任务。许多用户每天都会使用 Kimi 智能助手来总结文章，包括学术论文，以及一些第三方应用基于 Kimi 大模型开发的论文摘要工具。在这些实验中，我们比较了 Mooncake 架构和单机 vLLM 架构在处理 ArXiv 数据集时的性能。实验结果显示，Mooncake 在处理这些摘要任务时相较于单机 vLLM 有明显的优势。
上下文缓存的应用
Context Caching 基本原理
在传统的 LLM 交互中，每次用户与模型的对话都需要重新计算整个上下文，这包括了所有的历史信息和对话内容。这种方法意味着每次交互都需要为整个上下文支付计算费用，无论上下文中有多少信息是重复的。
上下文缓存的核心改进在于引入了一个“公共上下文”的概念。这个公共上下文包含了对话中不变的部分，比如背景信息或常见问题。通过缓存这个公共前缀，我们只需要为其支付一次计算费用，而不必在每次交互时重复支付。这样，每次交互的成本就大大降低了，用户只需要为每次的增量输入（即新的对话内容）以及存储公共上下文的费用付费。
Context Caching 使用流程
使用 Context Caching 的流程非常简洁明了，主要分为以下几个步骤：
创建缓存：首先，你需要创建一个缓存实例。这个过程通常非常快速，大约需要 30 到 40 秒的时间。一旦缓存创建完成，它就可以被用于后续的交互。
使用缓存：创建缓存后，你可以直接在对话或应用中使用它。由于缓存已经包含了必要的上下文信息，因此可以避免重复计算，提高响应速度和效率。
为了帮助开发者更容易地实现上下文缓存，我们在官方 GitHub 上提供了一些示例代码，这些代码覆盖了多种编程语言，包括 Python、Node.js 等，以便开发者能够快速上手并集成到自己的项目中。
Context Caching 收费模式
我们的上下文缓存（Context Caching）技术的收费模式已经进行了优化和调整。现在，创建缓存的成本非常低，而且是一次性的费用。调用缓存的费用也几乎可以忽略不计，主要的成本瓶颈在于存储空间的费用。为了鼓励更多开发者使用这项技术，我们最近对价格进行了调整，降价幅度达到了 50%。原先的费用是 10 元，现在降低到了 5 元。我们希望通过这样的降价措施，能够激励开发者更广泛地采用上下文缓存技术。
在最近的“GOGC 黑客松”和之前的“Adventure X 黑客松”中，我们特别表彰了使用上下文缓存 API 最多的开发者，以此鼓励大家探索和利用这项技术。我们注意到，尽管上下文缓存技术能带来便利和性能优化，但仍有许多开发者没有充分利用它。例如，之前有一个非常受欢迎的应用，它在短时间内有大量的调用。如果该应用当时使用了上下文缓存技术，开发者本可以节省一大笔开支。实际上，我们与开发者进行调研后发现，由于大模型产生的费用过高，他们不得不调整其提示词，尽可能简化内容，甚至删除了一些原本设计的游戏玩法、规则和元素，这无疑是一种遗憾。如果现在有类似的爆款应用出现，它们完全可以利用上下文缓存技术来提升用户体验并有效控制成本。
Context Caching 应用技巧
我们的官方 API 小助手的线上实际调用情况可以通过下图进行展示。目前，我们大约有 20 个开发者社群，每个社群的每小时调用情况都能清晰地反映出来。从图中可以看出，在每天凌晨 2 点到早上 8 点这段时间，社群内几乎没有活动，因为大家都在休息。一旦进入白天的工作时段，尤其是在早上、中午以及晚上下班后的时间段，社群内的活动会显著增加，这也是我们的 API 小助手最为忙碌的时段。
对于不同的开发者来说，他们的应用类型可能会影响调用频率的曲线。例如，如果应用是工具类或生产力类的，那么在工作日的白天可能会有更多的调用。相反，如果应用更偏向于娱乐或游戏，那么在晚上和周末的调用可能会更加频繁。
具体到我们小助手的情况：
在上午 9 点时，我们需要在调用 / chat/ completions 接口时，在 Headers 中添加 x-Msh-context-Cache 以启用 Cache， 同时添加 X-Msh-context-Cache-Reset-TTLHeader 以更新 Cache 存活期，这里以小助手为例，存活期为 3600s，即 1 小时；
由于我们要在凌晨。点结束 Cache，因此夜间 23 点是我们最后一次刷新 Cache 存活期的时点，在此之后，我们需要移除 Headers 中的 x-Msh-context-cache-Reset-TTL 参数，以保证 Cache 能在 0 点被顺利移除；
以 Python 代码为例，大致的代码逻辑为：
我们具体分析了一天中的数据，发现通过在特定时间点，如上午 9 点和晚上 24 点，存储上下文缓存，可以显著降低费用消耗，大约能节省 3/4 的成本。这种策略对于那些像我们的 API 小助手这样需要频繁交互的应用场景尤其有效。如果你的应用也属于这种类型，那么上下文缓存技术将是一个非常值得尝试的优化手段。
目前，我们也观察到市场上其他一些解决方案，它们通过不同的方式处理上下文缓存的复杂性。一些系统选择将这些复杂操作隐藏起来，让用户无需手动管理缓存的存储和删除，而是将这些任务完全交给调度引擎，由它智能地进行资源分配和优化。我认为手动管理和自动调度各有优势，手动管理提供了更多的灵活性和控制权，而自动调度则简化了操作流程，降低了用户的使用门槛。未来，我们可能会看到这两种方法的融合，以满足不同用户的需求和偏好。
Context Caching 适用场景
上下文缓存技术的应用场景非常广泛，尤其适合那些频繁进行请求并且需要重复引用大量初始上下文信息的场景。在这些情况下，上下文缓存能够显著提升处理效率，主要表现在缩短首次 Token 生成的时间，同时大幅降低 Token 消耗的费用。在我们的线上生产环境中，随着 Mooncake 架构从最初的灰度测试到现在的全面部署，Kimi 智能助手能够每天处理的请求量增加了 75%。这也是为什么用户最近感觉到“Kimi 累了”的情况有所减少的原因之一。
参考资料
为了帮助大家更深入地了解上下文缓存及相关技术，我们附上了一些参考资料。特别推荐其中一篇由爱丁堡大学的傅瑶撰写的基础总结，该文详细介绍了长上下文大型模型推理优化技术，包括最新的进展和其他技术的对比分析，以及哪些技术可以组合使用等。
Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving
Data Engineering for Scaling Language Models to 128K Context
Large Lanquage Model Based Long Context Modeling Papers and Bloas
Full Stack Transformer Inference Optimization Season 2: Deploying Long-Context Models
演讲嘉宾介绍
唐飞虎，月之暗面高级研发工程师、开发者关系负责人。前谷歌工程师、ACM/ICPC 亚洲赛区金牌、微软编程之美挑战赛冠军、第一届万向实验室通证经济设计大赛冠军。
