1. 项目缘起当大模型推理遇上昇腾集群的“甜蜜烦恼”最近半年我几乎把所有时间都泡在了昇腾Ascend集群上折腾各种大模型的训练和推理。昇腾的算力确实猛尤其是对于国产化替代这条路它的生态和性能表现越来越扎实。但做久了就会发现训练其实相对“省心”有成熟的框架和模式真正让人头疼的是生产环境下的大规模、高并发、长序列的在线推理服务。想象一下这个场景你部署了一个千亿参数的大模型作为智能客服或者内容生成服务用户请求源源不断每个请求的对话历史Context可能长达几千甚至上万个token。这时候两个核心痛点会立刻浮出水面计算资源“空转”与浪费每个用户的独立请求模型都需要从头到尾计算一遍。即使两个用户的提问高度相似比如都问“今天天气如何”或者同一个用户连续提问模型也无法复用之前已经计算过的中间结果。大量的计算特别是注意力机制中的Key和Value矩阵计算也就是我们常说的KVCache被重复执行导致宝贵的昇腾NPU算力被低效消耗。内存墙与吞吐量瓶颈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 MooncakeKVCache内存的“精算师与共享公寓管家”如果说Kthena管的是“计算任务怎么跑”那么Mooncake管的就是“内存怎么用”。它的焦点完全集中在Transformer推理过程中最宝贵也最麻烦的资源——KVCache上。Mooncake要解决的核心问题是KVCache的“动态性”和“碎片化”动态性每个请求的序列长度是实时变化的无法预知。碎片化如果每个请求独占连续内存随着请求的创建和结束内存中会出现大量无法被新请求利用的“空洞”外部碎片。Mooncake的核心技术PagedAttention分页注意力的实现这是Mooncake的基石。它受操作系统虚拟内存分页管理的启发将KVCache在逻辑上划分为固定大小的“块”Block比如每块存储16或32个token的KV向量。物理上这些块是预先申请好的一大块连续设备内存池。块级内存管理当一个新请求到来时Mooncake不是直接为它分配一个可能很大的连续空间而是从内存池中分配若干个空闲的块给它。请求的KVCache在逻辑上是一个由这些块组成的链表。序列增长时就追加分配新的块请求结束时它占用的所有块被标记为空闲归还给内存池供其他请求使用。块共享与复用这是提升效率的关键。Mooncake可以识别不同请求之间的共享前缀。例如系统提示词System Prompt或者多个用户共用的知识库前缀它们的KVCache可以被计算一次然后以“只读”块的形式被多个请求引用。这直接避免了重复计算节省了大量算力。与计算框架协同Mooncake需要修改模型前向传播中的注意力计算部分。它提供了一个定制化的Attention算子这个算子知道如何根据请求的“块映射表”从分散的物理块中 gather 出正确的Key和Value向量进行计算。这意味着模型代码需要做一定的适配才能接入Mooncake。Mooncake带来的直接收益内存利用率大幅提升消除了外部碎片内存池化使利用率接近100%。并发量显著增加同样大小的显存可以容纳更多请求的KVCache。计算开销降低通过共享块重复的注意力计算被消除。支持超长序列由于内存是块式管理的理论上可以支持远超单卡显存容量的超长序列只要内存池足够大。2.3 二者如何协同工作理解了各自角色它们的协作流程就清晰了服务启动Kthena首先启动根据配置将模型加载到多张昇腾卡上并启动服务端点。请求接入用户请求到达Kthena的调度器。资源分配Kthena将请求信息如输入ids传递给Mooncake。Mooncake为该请求分配逻辑块并处理可能的共享块逻辑。分布式执行Kthena调度器将包含了输入数据和Mooncake块信息的计算任务下发到各个模型并行阶段所在的NPU上。定制化计算在每个NPU上模型的前向传播会调用Mooncake提供的Attention算子。该算子根据当前请求的块信息从Mooncake管理的全局内存池中取得正确的KV数据完成注意力计算并更新KVCache写入新的块。生成与循环生成一个token后流程重复直到生成结束。请求完成后Mooncake回收其占用的所有块。这个架构下Kthena和Mooncake是松耦合的。理论上你可以将Mooncake与其他分布式框架结合或者将Kthena用于不需要KVCache复用的场景。但它们的组合确实为昇腾上的大模型推理提供了一套完整的“算力内存”优化方案。3. 昇腾环境下的部署实战从零到一的踩坑记录理论很美好但部署过程才是真正的试金石。我们的测试环境是8台搭载4颗昇腾910B NPU的服务器共32卡操作系统为CentOS 7.6驱动和固件版本为配套的最新版本。以下是我从环境准备到成功跑通Demo的完整步骤和关键坑点。3.1 基础环境准备避开版本“雷区”这是最基础也最容易出问题的一步。昇腾的软件栈包括驱动、固件、CANN异构计算架构以及AI框架MindSpore/PyTorch。核心教训务必严格对照官方发布的版本配套表社区版本迭代快用错版本组合会导致各种离奇错误。驱动与固件从华为昇腾社区下载对应服务器型号和OS的驱动包。安装后使用npu-smi info命令确认所有NPU状态正常。这里常遇到的问题是内核版本不匹配如果遇到可能需要升级系统内核或寻找对应内核版本的驱动。CANN工具包这是昇腾计算的基础软件层。我们选择了CANN 7.0.RC1版本因为它对后续要安装的PyTorch适配版本支持较好。安装时注意设置好环境变量特别是ASCEND_HOME和LD_LIBRARY_PATH。AI框架选择Kthena和Mooncake对MindSpore原生支持更好但我们的模型基于PyTorch因此选择了PyTorch 2.1 昇腾适配插件torch_npu的路线。这里有个大坑PyTorch、torch_npu、CANN三者的版本必须精确匹配。我们最终采用的组合是torch2.1.0torch_npu2.1.0.post3CANN7.0.RC1。安装torch_npu后务必执行python3 -m torch_npu.test.test_npu进行基础功能测试。3.2 Kthena的编译与安装攻克依赖难关Kthena的源码通常需要从GitHub仓库克隆并编译。它的依赖项较多包括gRPC、protobuf、abseil-cpp等。关键步骤克隆代码切换到与你的CANN和PyTorch版本匹配的分支。重点修改CMakeLists.txt或编译脚本中的昇腾相关路径。你需要明确指定CANN的安装路径、HCCL库的路径等。例如-DASCEND_PATH/usr/local/Ascend/ascend-toolkit/latest -DHCCL_PATH/usr/local/Ascend/ascend-toolkit/latest/lib64编译过程中最常见的错误是“找不到xxx.h”或“undefined reference to”。这几乎都是因为依赖库路径不对或版本冲突。我们的解决方案是在服务器上使用conda创建一个干净的Python环境并在该环境中用pip安装所有Python依赖如grpcio对于C依赖尽量使用系统包管理器yum安装开发版并确保CMake能正确找到它们。编译成功后你会得到kthena_server可执行文件以及Python客户端库。将Python包安装到你的环境pip install -e ./python。3.3 Mooncake的集成模型适配是关键Mooncake通常以Python库的形式提供但其核心包含需要编译的C/CUDA对于昇腾是NPU算子。获取与编译克隆Mooncake仓库。它的编译同样需要指向昇腾的CANN路径。重点关注其setup.py或build.py脚本确保--npu选项被正确启用并且编译器能找到ascendcl等头文件和库。模型改造这是工作量最大的部分。Mooncake无法直接作用于原始的PyTorchnn.MultiheadAttention或F.scaled_dot_product_attention。你需要找到模型中所有自注意力Self-Attention和解码器交叉注意力Cross-Attention的模块。将这些模块替换为Mooncake提供的对应模块例如moondream.Attention。这些模块的构造函数和forward函数的签名发生了变化需要传入Mooncake的BlockManager实例和请求的block_table等信息。你需要编写一个逻辑为每个请求维护一个block_table块表记录该请求的KVCache分布在哪些物理块上。这个表会在生成过程中动态更新并传递给Attention模块。初始化内存池在服务启动时需要根据你的显存大小和预期并发数初始化Mooncake的内存池确定块大小block_size和总块数。这是一个权衡艺术块太小管理开销大块太大内部碎片多。我们经过测试对于13B-70B的模型16或32个token每块是一个不错的起点。3.4 配置文件与启动串联整个系统现在你需要编写配置文件把Kthena、你的模型和Mooncake串起来。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] ... # 以此类推启动Kthena服务./kthena_server --config config.yaml --model-path ./model_weights --http-port 8080服务启动后会加载模型到各张NPU并监听端口。客户端请求使用Kthena的Python客户端或直接发送HTTP请求。请求体中除了常规的prompt和generation_config现在还需要包含一个由Mooncake客户端生成的cache_id或block_table信息。Mooncake的服务端部分BlockManager需要与你的模型代码运行在同一个进程内通常作为全局单例存在。4. 性能调优与效果验证数据不会说谎部署成功只是第一步真正的价值在于性能提升。我们设计了一系列对比实验。测试模型LLaMA-13B测试场景基准组使用原生PyTorch torch_npu单请求无KVCache复用。实验组AKthena分布式4卡流水线并行无Mooncake即每个请求独立KVCache。实验组BKthena分布式 Mooncake启用KVCache共享共享系统提示词。测试负载长上下文问答输入序列长度固定为4096 token输出长度128 token。模拟文档分析场景。多轮对话模拟模拟10个并发的用户对话每个用户进行5轮问答每轮输入约100 token输出50 token。对话间共享一部分系统指令。关键性能指标吞吐量Tokens/s整个集群每秒处理的输入输出token总数。首Token延迟TTFT从请求发出到收到第一个输出token的时间。内存占用通过npu-smi监控NPU的HBM使用率。测试结果与分析测试场景配置方案平均吞吐量 (Tokens/s)平均TTFT (ms)峰值HBM占用 (GB/卡)支持最大并发数长上下文 (4096128)基准组 (单卡)451850281长上下文 (4096128)实验组A (4卡无复用)1622100~26 (每卡)4长上下文 (4096128)实验组B (4卡有复用)1952050~22 (每卡)10多轮对话模拟 (10并发)基准组 (单卡)崩溃 (OOM)---多轮对话模拟 (10并发)实验组A (4卡无复用)280波动大~28 (每卡)10 (已满)多轮对话模拟 (10并发)实验组B (4卡有复用)520稳定~250~24 (每卡)可扩展至20结论非常明显分布式提升算力利用率实验组A对比基准组吞吐量提升约3.6倍162/45略低于理想的4倍这是由于流水线并行存在气泡Bubble开销。这证明了Kthena分布式推理的有效性。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是服务端状态。我们需要一个机制将客户端的会话IDSession 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引入我们的昇腾推理栈虽然前期投入了相当的工程精力但带来的收益是决定性的。它不仅仅是一个性能优化工具更是一种面向大模型推理服务的架构范式转变——从“单次请求、独占资源”的粗放模式转向“持续服务、资源共享”的精细化运营模式。对于后来者我的核心建议是明确需求如果你的场景是高并发、对话式、且请求间存在共享前缀如系统提示、知识库那么这套方案的收益会非常大。如果是简单的单次文本补全收益可能不那么明显。循序渐进不要试图一步到位。先从单卡、小模型、禁用Mooncake开始确保Kthena的基础分布式推理能跑通。然后引入Mooncake测试KVCache复用的正确性。最后再上大规模集群和真实流量。强化可观测性在集成之初就要投入资源构建完善的日志、指标和追踪系统。这是你后期排查问题、进行性能调优的“眼睛”。关注社区Kthena和Mooncake都是活跃的开源项目版本迭代很快。密切关注其Releases和Issues很多你遇到的坑可能已经有人踩过并修复了。这次实践让我们看到在大模型推理这条赛道上软件栈的创新与硬件算力的提升同等重要。Kthena和Mooncake这样的项目正是在填补从“模型能跑”到“服务高效、经济”之间的关键鸿沟。随着模型规模的进一步扩大和应用场景的深化这种专注于推理效率的底层系统优化其价值只会越来越凸显。