昇腾大模型推理优化：Kthena与Mooncake实现KVCache复用与分布式加速

1. 项目缘起：当大模型推理遇上昇腾集群的“甜蜜烦恼”

最近半年，我几乎把所有时间都泡在了昇腾（Ascend）集群上，折腾各种大模型的训练和推理。昇腾的算力确实猛，尤其是对于国产化替代这条路，它的生态和性能表现越来越扎实。但做久了就会发现，训练其实相对“省心”，有成熟的框架和模式；真正让人头疼的，是生产环境下的大规模、高并发、长序列的在线推理服务。

想象一下这个场景：你部署了一个千亿参数的大模型作为智能客服或者内容生成服务，用户请求源源不断，每个请求的对话历史（Context）可能长达几千甚至上万个token。这时候，两个核心痛点会立刻浮出水面：

1. 计算资源“空转”与浪费：每个用户的独立请求，模型都需要从头到尾计算一遍。即使两个用户的提问高度相似（比如都问“今天天气如何”），或者同一个用户连续提问，模型也无法复用之前已经计算过的中间结果。大量的计算，特别是注意力机制中的Key和Value矩阵计算（也就是我们常说的KVCache），被重复执行，导致宝贵的昇腾NPU算力被低效消耗。
2. 内存墙与吞吐量瓶颈：KVCache是Transformer解码（生成）过程中，为加速自回归生成而缓存的历史Key和Value向量。序列越长，KVCache占用的显存（在昇腾上叫HBM）就越大。在分布式推理中，如果每个请求独占一份KVCache，集群的可用内存会迅速被撑满，严重限制同时服务的请求数量（也就是吞吐量）。你加再多卡，也可能被内存限制卡住脖子。

这就像一家很火的餐厅（昇腾集群），厨师（NPU）炒菜速度很快，但每个顾客（请求）都必须从切菜开始完全独立做一份，后厨（显存）堆满了重复的半成品食材（KVCache），导致翻台率（吞吐量）根本上不去。

我一直在寻找能解决这两个痛点的方案。直到最近，在开源社区里看到了Kthena和Mooncake这两个项目的组合，尝试着在我们的昇腾910集群上部署和测试了一番，效果可以说是“柳暗花明”。Kthena × Mooncake，本质上是一套针对昇腾平台优化的、支持分布式推理与KVCache高效复用的推理服务引擎。它不是为了替代PyTorch或MindSpore，而是在它们之上，构建了一个更贴近生产部署场景的“服务层”。

简单来说，Kthena更像一个分布式推理的运行时调度与执行框架，它负责把一个大模型合理地切分到多张昇腾卡上（模型并行），并管理请求的生命周期。而Mooncake则是一个专注于KVCache内存管理与复用策略的库，它实现了诸如PagedAttention（分页注意力）等先进的内存管理技术，让不同请求可以安全、高效地共享KVCache内存块。

它们的结合，目标直指昇腾集群推理的终极效率：用更少的内存，服务更多的并发请求，榨干每一分NPU算力。接下来，我就结合实际的部署、测试和踩坑经历，详细拆解这套方案的核心原理、实操步骤以及那些官方文档里不会写的细节。

2. 核心组件拆解：Kthena与Mooncake各自扮演什么角色？

在深入实操之前，必须把Kthena和Mooncake的分工搞清楚。很多人容易把它们混为一谈，或者认为其中一个包含了另一个。实际上，它们是解耦的、各司其职的组件，通过清晰的接口进行协作。

2.1 Kthena：昇腾分布式推理的“交通指挥官”

你可以把Kthena理解为一个专为昇腾NPU设计的、轻量级分布式服务框架。它的核心职责不是实现模型算法，而是解决“如何让一个大规模模型在多个NPU上协同工作，并高效处理海量用户请求”这个系统性问题。

它的几个关键设计点：

• 模型并行抽象：Kthena提供了一套API，让你能够以相对直观的方式描述一个超大模型应该如何被切割到不同的设备上。例如，对于一个Transformer模型，你可以指定哪些层放在Device 0，哪些放在Device 1。它底层会处理好设备间的通信（如All-Reduce， All-Gather），这些通信原语都针对昇腾的HCCL（华为集合通信库）做了深度优化。
• 请求调度与流水线：面对蜂拥而至的请求，Kthena实现了动态的调度策略。它可以将不同的请求，甚至同一个请求生成的不同token，以流水线（Pipeline）的方式在多个阶段（Stage，对应模型的不同部分）上执行。这能有效提升NPU的利用率，避免某个设备空闲等待。
• 与推理引擎解耦：Kthena并不绑定某个特定的模型实现或推理引擎。它目前主要与MindSpore和PyTorch（通过昇腾适配插件）对接。你只需要用这些框架定义好你的模型，Kthena负责把这个模型“分布式化”并运行起来。
• 服务化接口：它通常提供gRPC或HTTP等标准网络接口，方便集成到现有的微服务架构中。你的前端应用只需要像调用普通API一样发送请求，无需关心背后模型分布在多少张卡上。

为什么需要Kthena？如果没有它，你在昇腾上做分布式推理，可能需要手动写大量的设备放置代码、通信同步逻辑和请求队列管理，复杂度极高且容易出错。Kthena把这些脏活累活封装了。

2.2 Mooncake：KVCache内存的“精算师与共享公寓管家”

如果说Kthena管的是“计算任务怎么跑”，那么Mooncake管的就是“内存怎么用”。它的焦点完全集中在Transformer推理过程中最宝贵也最麻烦的资源——KVCache上。

Mooncake要解决的核心问题是KVCache的“动态性”和“碎片化”：

• 动态性：每个请求的序列长度是实时变化的，无法预知。
• 碎片化：如果每个请求独占连续内存，随着请求的创建和结束，内存中会出现大量无法被新请求利用的“空洞”（外部碎片）。

Mooncake的核心技术：

1. PagedAttention（分页注意力）的实现：这是Mooncake的基石。它受操作系统虚拟内存分页管理的启发，将KVCache在逻辑上划分为固定大小的“块”（Block），比如每块存储16或32个token的KV向量。物理上，这些块是预先申请好的一大块连续设备内存池。
2. 块级内存管理：当一个新请求到来时，Mooncake不是直接为它分配一个可能很大的连续空间，而是从内存池中分配若干个空闲的块给它。请求的KVCache在逻辑上是一个由这些块组成的链表。序列增长时，就追加分配新的块；请求结束时，它占用的所有块被标记为空闲，归还给内存池，供其他请求使用。
3. 块共享与复用：这是提升效率的关键。Mooncake可以识别不同请求之间的共享前缀。例如，系统提示词（System Prompt）或者多个用户共用的知识库前缀，它们的KVCache可以被计算一次，然后以“只读”块的形式被多个请求引用。这直接避免了重复计算，节省了大量算力。
4. 与计算框架协同：Mooncake需要修改模型前向传播中的注意力计算部分。它提供了一个定制化的Attention算子，这个算子知道如何根据请求的“块映射表”，从分散的物理块中 gather 出正确的Key和Value向量进行计算。这意味着模型代码需要做一定的适配才能接入Mooncake。

Mooncake带来的直接收益：

• 内存利用率大幅提升：消除了外部碎片，内存池化使利用率接近100%。
• 并发量显著增加：同样大小的显存，可以容纳更多请求的KVCache。
• 计算开销降低：通过共享块，重复的注意力计算被消除。
• 支持超长序列：由于内存是块式管理的，理论上可以支持远超单卡显存容量的超长序列（只要内存池足够大）。

2.3 二者如何协同工作？

理解了各自角色，它们的协作流程就清晰了：

1. 服务启动：Kthena首先启动，根据配置将模型加载到多张昇腾卡上，并启动服务端点。
2. 请求接入：用户请求到达Kthena的调度器。
3. 资源分配：Kthena将请求信息（如输入ids）传递给Mooncake。Mooncake为该请求分配逻辑块，并处理可能的共享块逻辑。
4. 分布式执行：Kthena调度器将包含了输入数据和Mooncake块信息的计算任务，下发到各个模型并行阶段所在的NPU上。
5. 定制化计算：在每个NPU上，模型的前向传播会调用Mooncake提供的Attention算子。该算子根据当前请求的块信息，从Mooncake管理的全局内存池中取得正确的KV数据，完成注意力计算，并更新KVCache（写入新的块）。
6. 生成与循环：生成一个token后，流程重复，直到生成结束。请求完成后，Mooncake回收其占用的所有块。

这个架构下，Kthena和Mooncake是松耦合的。理论上，你可以将Mooncake与其他分布式框架结合，或者将Kthena用于不需要KVCache复用的场景。但它们的组合，确实为昇腾上的大模型推理提供了一套完整的“算力+内存”优化方案。

3. 昇腾环境下的部署实战：从零到一的踩坑记录

理论很美好，但部署过程才是真正的试金石。我们的测试环境是8台搭载4颗昇腾910B NPU的服务器（共32卡），操作系统为CentOS 7.6，驱动和固件版本为配套的最新版本。以下是我从环境准备到成功跑通Demo的完整步骤和关键坑点。

3.1 基础环境准备：避开版本“雷区”

这是最基础也最容易出问题的一步。昇腾的软件栈包括驱动、固件、CANN（异构计算架构）以及AI框架（MindSpore/PyTorch）。

核心教训：务必严格对照官方发布的版本配套表！社区版本迭代快，用错版本组合会导致各种离奇错误。

1. 驱动与固件：从华为昇腾社区下载对应服务器型号和OS的驱动包。安装后，使用npu-smi info命令确认所有NPU状态正常。这里常遇到的问题是内核版本不匹配，如果遇到，可能需要升级系统内核或寻找对应内核版本的驱动。
2. CANN工具包：这是昇腾计算的基础软件层。我们选择了CANN 7.0.RC1版本，因为它对后续要安装的PyTorch适配版本支持较好。安装时注意设置好环境变量，特别是ASCEND_HOME和LD_LIBRARY_PATH。
3. AI框架选择：Kthena和Mooncake对MindSpore原生支持更好，但我们的模型基于PyTorch，因此选择了PyTorch 2.1 + 昇腾适配插件（torch_npu）的路线。这里有个大坑：PyTorch、torch_npu、CANN三者的版本必须精确匹配。我们最终采用的组合是torch==2.1.0+torch_npu==2.1.0.post3+CANN==7.0.RC1。安装torch_npu后，务必执行python3 -m torch_npu.test.test_npu进行基础功能测试。

3.2 Kthena的编译与安装：攻克依赖难关

Kthena的源码通常需要从GitHub仓库克隆并编译。它的依赖项较多，包括gRPC、protobuf、abseil-cpp等。

关键步骤：

1. 克隆代码，切换到与你的CANN和PyTorch版本匹配的分支。
2. 重点修改CMakeLists.txt或编译脚本中的昇腾相关路径。你需要明确指定CANN的安装路径、HCCL库的路径等。例如：

   -DASCEND_PATH=/usr/local/Ascend/ascend-toolkit/latest -DHCCL_PATH=/usr/local/Ascend/ascend-toolkit/latest/lib64

3. 编译过程中最常见的错误是“找不到xxx.h”或“undefined reference to”。这几乎都是因为依赖库路径不对或版本冲突。我们的解决方案是：在服务器上使用conda创建一个干净的Python环境，并在该环境中用pip安装所有Python依赖（如grpcio）；对于C++依赖，尽量使用系统包管理器（yum）安装开发版，并确保CMake能正确找到它们。
4. 编译成功后，你会得到kthena_server可执行文件以及Python客户端库。将Python包安装到你的环境：pip install -e ./python。

3.3 Mooncake的集成：模型适配是关键

Mooncake通常以Python库的形式提供，但其核心包含需要编译的C++/CUDA（对于昇腾是NPU）算子。

1. 获取与编译：克隆Mooncake仓库。它的编译同样需要指向昇腾的CANN路径。重点关注其setup.py或build.py脚本，确保--npu选项被正确启用，并且编译器能找到ascendcl等头文件和库。
2. 模型改造：这是工作量最大的部分。Mooncake无法直接作用于原始的PyTorchnn.MultiheadAttention或F.scaled_dot_product_attention。你需要：
   • 找到模型中所有自注意力（Self-Attention）和解码器交叉注意力（Cross-Attention）的模块。
   • 将这些模块替换为Mooncake提供的对应模块，例如moondream.Attention。这些模块的构造函数和forward函数的签名发生了变化，需要传入Mooncake的BlockManager实例和请求的block_table等信息。
   • 你需要编写一个逻辑，为每个请求维护一个block_table（块表），记录该请求的KVCache分布在哪些物理块上。这个表会在生成过程中动态更新，并传递给Attention模块。
3. 初始化内存池：在服务启动时，需要根据你的显存大小和预期并发数，初始化Mooncake的内存池，确定块大小（block_size）和总块数。这是一个权衡艺术：块太小，管理开销大；块太大，内部碎片多。我们经过测试，对于13B-70B的模型，16或32个token每块是一个不错的起点。

3.4 配置文件与启动：串联整个系统

现在，你需要编写配置文件，把Kthena、你的模型和Mooncake串起来。

1. Kthena模型配置文件：这是一个JSON或YAML文件，描述模型并行策略。例如，对于一个40层的模型，在4卡上采用流水线并行（Pipeline Parallelism），你可以配置每个“阶段”（stage）负责10层。同时，你还需要指定模型文件的路径、tokenizer路径等。

   model_name: "my_llama_13b" tensor_parallel_size: 1 # 张量并行，我们先用流水线并行 pipeline_parallel_size: 4 # 4个流水线阶段 pipeline_parallel_stages: - device_ids: [0] # stage0 在卡0上，运行第0-9层 model_layer_range: [0, 10] - device_ids: [1] # stage1 在卡1上，运行第10-19层 model_layer_range: [10, 20] ... # 以此类推

2. 启动Kthena服务：

   ./kthena_server --config config.yaml --model-path ./model_weights --http-port 8080

   服务启动后，会加载模型到各张NPU，并监听端口。
3. 客户端请求：使用Kthena的Python客户端或直接发送HTTP请求。请求体中除了常规的prompt和generation_config，现在还需要包含一个由Mooncake客户端生成的cache_id或block_table信息。Mooncake的服务端部分（BlockManager）需要与你的模型代码运行在同一个进程内，通常作为全局单例存在。

4. 性能调优与效果验证：数据不会说谎

部署成功只是第一步，真正的价值在于性能提升。我们设计了一系列对比实验。

测试模型：LLaMA-13B测试场景：

• 基准组：使用原生PyTorch + torch_npu，单请求，无KVCache复用。
• 实验组A：Kthena分布式（4卡流水线并行），无Mooncake（即每个请求独立KVCache）。
• 实验组B：Kthena分布式 + Mooncake（启用KVCache共享，共享系统提示词）。

测试负载：

1. 长上下文问答：输入序列长度固定为4096 token，输出长度128 token。模拟文档分析场景。
2. 多轮对话模拟：模拟10个并发的用户对话，每个用户进行5轮问答，每轮输入约100 token，输出50 token。对话间共享一部分系统指令。

关键性能指标：

• 吞吐量（Tokens/s）：整个集群每秒处理的输入+输出token总数。
• 首Token延迟（TTFT）：从请求发出到收到第一个输出token的时间。
• 内存占用：通过npu-smi监控NPU的HBM使用率。

测试结果与分析：

结论非常明显：

1. 分布式提升算力利用率：实验组A对比基准组，吞吐量提升约3.6倍（162/45），略低于理想的4倍，这是由于流水线并行存在气泡（Bubble）开销。这证明了Kthena分布式推理的有效性。
2. KVCache复用带来质变：
   • 吞吐量飞跃：在并发场景下（实验组B vs A），吞吐量从280提升至520，提升近86%。这主要归功于Mooncake通过共享块避免了大量重复计算。
   • 内存效率大幅提升：峰值HBM占用降低了约15%（24GB vs 28GB），这使得在同等硬件下，系统能支撑的并发请求数几乎翻倍。这是PagedAttention内存池化消除碎片的效果。
   • 延迟更稳定：实验组A在并发时TTFT波动大，因为请求间会争抢内存带宽和计算资源。实验组B由于内存管理更高效，资源争抢减少，延迟更加平稳可控。

调优心得：Mooncake的块大小（block_size）对性能影响显著。我们最初设置为64，发现对于短序列（几十个token）的请求，内部碎片浪费严重。调整为16后，短序列请求的内存效率提升，但管理开销（块表更大）略有增加。需要根据你的实际请求长度分布进行压测和调整。一个实用的技巧是对不同长度的请求使用不同的块大小策略，但这需要更复杂的Mooncake定制。

5. 生产落地中的挑战与应对策略

在测试环境跑通后，我们尝试将其推向一个接近生产的环境，遇到了更多工程上的挑战。

5.1 故障排查与监控：让系统变得可观测

分布式系统加上复杂的内存管理，一旦出问题，日志就像天书。

• 挑战1：请求卡死或无响应。可能原因：某个NPU上的进程异常退出；Kthena调度器死锁；Mooncake块分配出现逻辑错误（如循环引用）。
  • 应对：我们为Kthena服务添加了详细的结构化日志，记录每个请求进入每个流水线阶段的时间。同时，利用昇腾的Ascend Profiler工具，定期采样NPU上的算子执行情况。对于Mooncake，我们增加了块分配/释放的审计日志，确保每个块的生命周期可追溯。
• 挑战2：内存缓慢增长直至OOM。这是最棘手的问题，通常是内存泄漏。
  • 应对：我们编写了监控脚本，定期通过Mooncake的API查询内存池的状态：总块数、已分配块数、每个请求持有的块列表。同时，监控NPU的HBM使用情况。一旦发现“已分配块数”只增不减，或者某个已结束请求的块未被释放，就能快速定位泄漏点。最终发现是我们自定义的请求上下文管理器中，在异常处理分支里忘记调用释放函数。
• 挑战3：性能抖动。在长时间运行后，吞吐量偶尔会突然下降。
  • 应对：通过监控发现，这与Linux系统的透明大页（THP）或NPU驱动内部的内存回收机制有关。我们调整了系统的THP设置为madvise，并为昇腾驱动设置了更积极的内存预留参数，减少了抖动。

5.2 与现有系统的集成：不是孤岛

我们的在线服务已有基于Flask的API网关和Redis缓存。

• 集成模式：我们没有直接用Kthena的HTTP端口对外暴露，而是保留原有的API网关。网关接收请求后，先进行鉴权、限流、输入校验等。然后，网关服务内部作为一个Kthena的gRPC客户端，将请求转发给Kthena集群。这样既利用了现有网关的能力，又隔离了内部推理服务的复杂性。
• 会话状态管理：Mooncake的KVCache是服务端状态。我们需要一个机制，将客户端的会话ID（Session ID）映射到Mooncake内部的cache_id和block_table。我们利用Redis来存储这个映射关系。当用户发起续聊请求时，网关从Redis中取出对应的cache_id发给Kthena，从而实现跨请求的KVCache持久化，真正实现“记住对话历史”。
• 优雅扩缩容：当需要增加推理节点时，Kthena支持动态添加新的“工作者”（Worker）。但Mooncake的内存池是每个进程独立的。我们的策略是：新启动的服务实例拥有独立的内存池，通过网关层的负载均衡器（如Nginx）将新会话导向新实例。对于老实例，等待其现有会话自然结束后再下线。这避免了状态迁移的复杂性。

5.3 对模型架构的假设与限制

Kthena和Mooncake这套方案并非银弹，它对模型架构有一些隐含假设：

• 仅限自回归Decoder模型：这套优化主要针对GPT、LLaMA等采用自回归生成方式的Decoder-only模型。对于Encoder-Decoder模型（如T5）或非自回归模型，KVCache的复用模式和价值可能不同。
• 注意力机制需标准：Mooncake的块管理逻辑紧密耦合标准的Transformer注意力计算。如果模型使用了高度定制化的、结构迥异的注意力变体（例如某些线性注意力），可能需要大幅修改Mooncake的算子才能适配。
• 量化与MOE的支持：我们测试了INT8量化模型，Kthena可以正常加载运行，但Mooncake需要确保其内存池管理和注意力计算能与量化后的数据格式兼容。对于Mixture of Experts (MOE) 模型，情况更复杂，因为不同token可能激活不同的专家，KVCache的共享逻辑需要专家路由信息的参与，目前社区还在探索中。

6. 总结与展望：一次有价值的架构升级

回顾整个从调研、部署、测试到集成的过程，将Kthena与Mooncake引入我们的昇腾推理栈，虽然前期投入了相当的工程精力，但带来的收益是决定性的。它不仅仅是一个性能优化工具，更是一种面向大模型推理服务的架构范式转变——从“单次请求、独占资源”的粗放模式，转向“持续服务、资源共享”的精细化运营模式。

对于后来者，我的核心建议是：

1. 明确需求：如果你的场景是高并发、对话式、且请求间存在共享前缀（如系统提示、知识库），那么这套方案的收益会非常大。如果是简单的单次文本补全，收益可能不那么明显。
2. 循序渐进：不要试图一步到位。先从单卡、小模型、禁用Mooncake开始，确保Kthena的基础分布式推理能跑通。然后引入Mooncake，测试KVCache复用的正确性。最后再上大规模集群和真实流量。
3. 强化可观测性：在集成之初，就要投入资源构建完善的日志、指标和追踪系统。这是你后期排查问题、进行性能调优的“眼睛”。
4. 关注社区：Kthena和Mooncake都是活跃的开源项目，版本迭代很快。密切关注其Releases和Issues，很多你遇到的坑可能已经有人踩过并修复了。

这次实践让我们看到，在大模型推理这条赛道上，软件栈的创新与硬件算力的提升同等重要。Kthena和Mooncake这样的项目，正是在填补从“模型能跑”到“服务高效、经济”之间的关键鸿沟。随着模型规模的进一步扩大和应用场景的深化，这种专注于推理效率的底层系统优化，其价值只会越来越凸显。