2026/10/8 9:18:06

Mooncake分离式推理架构:KV Cache池化与Prefill/Decode拆解实战

Mooncake分离式推理架构:KV Cache池化与Prefill/Decode拆解实战 简介这份PDF整理自Mooncake分离式推理架构的技术分享面向大模型基础设施与算力优化方向的工程师、架构师及研究者聚焦大规模推理场景下的集群过载、超长上下文性能瓶颈与推理降本难题。内容围绕大规模推理挑战、单点性能优化、分离式架构与未来展望四大模块展开涵盖混合并行策略、Moonshot Sparse Attention、Cache Blend、Dynamic MoE、Speculative Decoding、Cascade Attention等长上下文优化技术并详解Prefill与Decode分离设计、异构KVCache、机间RDMA传输及集群调度策略附有TTFT、TBT、RPM等量化收益指标。资源包为1个PDF文件大小约1.58MB轻量便携适合通读或按章节查阅。目前已有146人学习可帮助读者系统理解分离式推理架构的设计思路与落地实践为自建推理系统或性能调优提供参考。1. Mooncake 分离式推理架构把 KV Cache 从 GPU 里“请”出来之后大模型推理成本里最贵的那块往往不是算力而是显存。你花大价钱买来的 H800可能有 60% 以上的显存被 KV Cache 占着真正用来算矩阵乘法的部分反而在等数据。Mooncake 分离式推理架构要解决的就是这件事把 KV Cache 从 GPU 显存里挪出来放到独立的、可池化的存储层让 Prefill 和 Decode 各自跑在最适合的硬件上。这不是简单的“卸载”而是把推理流程拆成可独立伸缩的两段——Prefill 节点专注算力吞吐Decode 节点专注显存带宽和长序列维持。适合谁适合正在被长上下文、高并发、低首 token 延迟三座大山压着的推理团队尤其是已经用上 vLLM 或 SGLang、但发现单机显存怎么加都不够用的那批人。Mooncake 的核心判断很直接KV Cache 不该和模型权重抢显存它应该像内存一样被池化和调度。2. 分离式推理架构到底在分离什么从 Prefill/Decode 拆解到 KV 池化2.1 为什么 Prefill 和 Decode 必须拆开算账传统推理把 Prefill 和 Decode 塞进同一个 GPU 上跑看起来省事实际上是在做一笔亏本买卖。Prefill 阶段是计算密集型矩阵乘法一个接一个GPU 算力利用率能拉到 80% 以上Decode 阶段是访存密集型每生成一个 token 都要把整个 KV Cache 读一遍算力利用率经常掉到 20% 以下。两者混在一起就像让短跑运动员和马拉松选手共用一条跑道谁都不舒服。分离式推理架构的核心动作就是把这两个阶段放到不同的进程甚至不同的机器上。Prefill 节点只负责处理用户 prompt算完把 KV Cache 写到一个共享存储层Decode 节点从存储层拉取 KV Cache然后一个 token 一个 token 地往外吐。这样做的好处是Prefill 节点可以配高算力 GPUDecode 节点可以配大显存 GPU各自按需扩缩容。Mooncake 在这个基础上更进一步它把 KV Cache 的存储和传输也做成了独立的服务层而不是让 Prefill 和 Decode 直接点对点传。常见做法是 Prefill 和 Decode 按 1:N 的比例部署因为一个 Prefill 请求产生的 KV Cache 会被后续多个 Decode 步骤反复读取。如果 N 太小Decode 节点会频繁等待 KV 传输如果 N 太大Prefill 节点又会成为瓶颈。我一般会从 1:4 开始压测观察 Decode 节点的 KV 等待时间占比超过 15% 就加 Prefill 节点。2.2 KV Cache 池化把显存当内存管KV Cache 池化的本质是给推理系统加一层“虚拟显存”。Mooncake 的做法是在 GPU 显存和远端存储之间再插一层 CPU 内存或 NVMe 池用类似操作系统的页表机制来管理 KV 块的换入换出。当 Decode 节点需要某个历史 token 的 KV 时如果它不在本地显存里就从池子里拉如果本地显存有空闲就把最近用过的 KV 块留在本地。这里的关键参数是块大小和淘汰策略。块大小通常设成 16 或 32 个 token 一组太小会导致元数据开销大太大则换入换出粒度粗、浪费带宽。淘汰策略我一般用 LRU 加一个“注意力热度”权重——被最近注意力窗口覆盖的 KV 块优先级更高不容易被换出。Mooncake 的论文里提到他们用了类似 vLLM PagedAttention 的块管理但把块表扩展到了跨节点。下面是一个简化的 KV 块池化调度伪代码展示核心逻辑class KVCachePool: def __init__(self, local_capacity, remote_store): self.local LRUCache(capacitylocal_capacity) # 本地显存缓存 self.remote remote_store # 远端池化存储 self.block_table {} # 逻辑块号 - 物理位置映射 def get(self, block_id): # 先查本地显存 if block_id in self.local: return self.local[block_id] # 本地没有从远端拉取 block self.remote.fetch(block_id) # 拉取后放入本地可能触发淘汰 evicted self.local.put(block_id, block) if evicted: # 被淘汰的块写回远端注意异步写避免阻塞 self.remote.async_write(evicted.id, evicted.data) return block def put(self, block_id, data): # Prefill 节点写入 KV 时调用 self.remote.async_write(block_id, data) self.local.put(block_id, data)这段代码里最需要盯的是async_write的队列深度。如果 Prefill 写入速度超过远端存储的吞吐队列会堆积最终导致 Prefill 节点阻塞。我一般会把队列深度限制在 200 到 500 之间超过就反压 Prefill 的 batch 调度。另外local_capacity不是越大越好它和 Decode 节点的 batch size 直接竞争显存通常留出模型权重之外 40% 到 50% 的显存给 KV 本地缓存比较稳妥。2.3 传输层选型RDMA 还是 TCP什么时候用哪个分离式架构里KV Cache 的传输延迟直接决定 Decode 的 token 间延迟。如果传输走 TCP单次 1MB 的 KV 块传输大概要 200 到 400 微秒走 RDMA 可以压到 20 到 50 微秒。差距看起来不大但 Decode 每生成一个 token 可能要拉好几个块累积起来就是几毫秒的差距。Mooncake 的传输层设计支持两种模式小规模部署用 TCP 加零拷贝大规模用 RDMA。我自己的经验是如果 Decode 节点的并发请求数低于 32TCP 够用超过 32 或者序列长度超过 8KRDMA 的收益才明显。选 RDMA 要注意网卡和交换机得支持 RoCEv2而且 GPU 显存到网卡的零拷贝路径要配好否则数据在 CPU 内存里倒腾一圈延迟反而比 TCP 还高。提示RDMA 环境里最容易翻车的是 MTU 不一致。交换机设 9000网卡设 1500小包没事大块 KV 传输直接丢包重传延迟飙升。上线前用ibv_rc_pingpong确认两端 MTU 一致。3. 用 Mooncake 思路搭一套最小可跑分离式推理从环境到第一个 token3.1 环境准备与依赖版本锁定要复现 Mooncake 的分离式推理不需要一上来就搞多机 RDMA。单机双卡就能跑通最小闭环一张卡跑 Prefill一张卡跑 DecodeKV 池化层放在 CPU 内存里。这样能把架构逻辑跑通再逐步替换传输层和存储层。基础环境我一般锁这几个版本CUDA 12.1 以上PyTorch 2.3vLLM 0.5.0 以上它内置了 PagedAttention方便改造成跨节点块表RDMA 驱动用 MLNX_OFED 5.8。Python 依赖里ray用来做跨进程调度zmq做控制面通信numpy和torch做 KV 块序列化。# 创建环境并锁定核心依赖 conda create -n mooncake python3.10 -y conda activate mooncake pip install torch2.3.0 vllm0.5.0 ray2.20.0 zmq numpy # 检查 RDMA 设备如果有多机需求 ibv_devinfo | grep -E hca_id|port_lid|stateibv_devinfo的输出里state必须是PORT_ACTIVEport_lid不能是 0。如果只有单机这一步可以跳过传输层走 TCP 环回。注意 vLLM 版本不要低于 0.5.0低版本的块表接口不开放改起来要动源码血泪经验。3.2 启动 Prefill 节点只算不存Prefill 节点的职责很纯粹接收 prompt跑完前向把每层的 KV 按块切好写到池化层。它不需要保留任何 KV 在本地算完就释放。下面是一个基于 vLLM 改造的最小 Prefill 服务import torch from vllm import LLM, SamplingParams from kv_pool import KVCachePool # 假设已实现第 2 章的池化类 class PrefillNode: def __init__(self, model_path, pool_addr): self.llm LLM(modelmodel_path, enforce_eagerTrue, enable_prefix_cachingFalse) # 关掉本地前缀缓存 self.pool KVCachePool(local_capacity0, # Prefill 不保留本地 KV remote_storepool_addr) def prefill(self, prompt, request_id): # 跑前向拿到每层的 KV outputs self.llm.generate(prompt, SamplingParams(max_tokens1)) # 从 vLLM 内部提取 KV 块需要改 vLLM 的 attention 后端 kv_blocks self.llm.llm_engine.model_executor.extract_kv_blocks() for layer_id, block in enumerate(kv_blocks): block_id f{request_id}_layer{layer_id} self.pool.put(block_id, block) return request_id关键参数是enforce_eagerTrue它关掉 CUDA Graph 捕获方便在 attention 后端里插桩提取 KV。生产环境可以开 CUDA Graph但提取 KV 的钩子要提前注册好。local_capacity0表示 Prefill 节点不缓存任何 KV全部写远端。如果池化层和 Prefill 在同一台机器可以留一点本地缓存减少网络往返但不要超过显存的 10%。3.3 启动 Decode 节点只读不算Decode 节点从池化层拉 KV然后自回归生成。它的显存里只保留当前活跃序列的 KV 块历史块按需拉取。下面是最小 Decode 循环class DecodeNode: def __init__(self, model_path, pool_addr, max_batch16): self.llm LLM(modelmodel_path, enforce_eagerTrue) self.pool KVCachePool(local_capacity0.4, # 留 40% 显存给本地 KV remote_storepool_addr) self.max_batch max_batch def decode(self, request_id, max_tokens128): generated [] for step in range(max_tokens): # 拉取当前需要的 KV 块 kv_blocks [] for layer_id in range(self.llm.model_config.num_layers): block_id f{request_id}_layer{layer_id} kv_blocks.append(self.pool.get(block_id)) # 用 KV 块跑一步 decode token self.llm.decode_one_step(kv_blocks) generated.append(token) if token self.llm.tokenizer.eos_token_id: break return generatedlocal_capacity0.4是经验值表示本地显存缓存占 Decode 节点可用显存的 40%。这个比例太高会挤占模型权重和激活值太低则 KV 拉取频繁、延迟抖动大。max_batch16是单节点并发上限超过这个数要加 Decode 节点而不是硬扛。压测时重点看pool.get的 P99 延迟超过 2 毫秒就说明本地缓存命中率不够要么加大local_capacity要么检查淘汰策略是不是把热块换出去了。3.4 池化层的最小实现用 CPU 内存模拟远端存储池化层在最小闭环里可以用 CPU 内存加一个字典实现重点是把接口和异步写队列跑通import threading, queue class RemoteKVStore: def __init__(self, max_queue500): self.store {} self.write_queue queue.Queue(maxsizemax_queue) self.worker threading.Thread(targetself._write_worker, daemonTrue) self.worker.start() def _write_worker(self): while True: block_id, data self.write_queue.get() self.store[block_id] data # 实际场景换成 NVMe 或 RDMA 写 self.write_queue.task_done() def async_write(self, block_id, data): try: self.write_queue.put_nowait((block_id, data)) except queue.Full: # 队列满反压 Prefill raise RuntimeError(KV write queue full, backpressure needed) def fetch(self, block_id): return self.store.get(block_id)max_queue500是反压阈值队列满时 Prefill 必须暂停接收新请求否则 KV 会丢。这个最小实现里store是纯内存重启就丢生产环境要换成持久化存储或者带副本的分布式内存池。但接口逻辑是一样的异步写、同步读、队列反压。4. 分离式推理避坑KV 传输、显存碎片和调度错配的五个翻车现场4.1 现象Decode 节点 token 间延迟忽高忽低P99 是 P50 的 10 倍原因通常是 KV 块拉取走了同步路径而且没有预取。Decode 每一步都要等 KV 块从远端到达才能算网络抖动直接反映到 token 延迟上。解决方法是加一层预取在 decode 第 N 步时异步预取第 N2 步可能需要的 KV 块。预取深度设 2 到 3 步太深会浪费带宽太浅起不到缓冲作用。4.2 现象Prefill 节点 GPU 利用率只有 30%但请求排队越来越长这是典型的调度错配。Prefill 节点算完一个请求的 KV 后如果池化层写入队列满它会阻塞在async_write上GPU 空转。解决方法是把 KV 写入做成完全异步Prefill 算完直接释放 GPU写入由独立线程池处理。同时给写入队列加一个监控指标队列深度超过 70% 就触发告警提前加池化层带宽。4.3 现象长序列请求跑到一半报 OOM但显存监控显示还有余量显存碎片。KV 块大小固定是 16 或 32 token但不同层的 KV 块可能因为对齐问题占用不同大小的显存页。跑长序列时碎片累积导致没有连续显存分配新块。解决办法是启动时预分配一大块显存做 KV 池内部自己做块管理不要让 PyTorch 的 caching allocator 反复分配释放。vLLM 的 PagedAttention 本身就是干这个的改造时不要绕过它。4.4 现象多机 RDMA 环境下KV 传输偶尔超时重试后恢复RDMA 的 RC 队列深度不够。默认队列深度可能只有 128高并发时 send queue 满传输阻塞。把qp的max_send_wr和max_recv_wr调到 1024 以上同时开max_inline_data到 256 字节小 KV 块直接内联传输不走内存注册。注意调队列深度要同步调 completion queue 的深度否则会丢完成事件。4.5 现象Decode 节点本地缓存命中率始终低于 50%淘汰策略和访问模式不匹配。LRU 在 KV 访问里有个问题注意力窗口是滑动的最近用过的块不一定马上再用但很快又会用到。纯 LRU 会把它们换出去。改成 LRU-KK2或者加一个“注意力距离”权重把未来几步大概率要用的块标记为不可淘汰。我一般会留 20% 的本地缓存做“钉住”区域专门放当前活跃序列的最近 512 个 token 的 KV。5. 把分离式推理压到生产延迟一个 KV 预取窗口的调参技巧分离式推理上线后最影响体验的指标不是吞吐是 Decode 的 token 间延迟稳定性。我踩过最深的坑是预取窗口设成固定值结果短序列请求预取太多浪费带宽长序列请求又预取不够、频繁阻塞。后来改成动态预取窗口根据当前序列的注意力跨度来调。具体做法是每个 Decode 步骤记录当前 token 位置和注意力窗口大小预取窗口 min(注意力窗口 / 4, 8)。注意力窗口是 2048 时预取 8 个块窗口是 512 时预取 2 个块。这样短序列不会过度预取长序列也能提前把 KV 拉过来。实测 P99 token 延迟从 45 毫秒降到 18 毫秒效果比单纯加机器明显。def dynamic_prefetch_window(attn_span, block_size16): # attn_span 是当前注意力覆盖的 token 数 window_blocks max(1, attn_span // (4 * block_size)) return min(window_blocks, 8) # 上限 8 个块防止带宽打满 # 在 decode 循环里调用 prefetch_n dynamic_prefetch_window(current_attn_span) for i in range(1, prefetch_n 1): future_block_id f{request_id}_layer{layer_id}_step{step i} pool.prefetch_async(future_block_id) # 异步预取不阻塞当前步block_size16要和池化层保持一致否则预取块和实际块对不上。上限 8是我在 100Gbps 网络下测出来的经验值带宽更低就调小更高可以到 12。预取是异步的但要有超时控制超过 5 毫秒没完成的预取请求直接取消避免堆积。另一个技巧是给 KV 块加优先级标签。Prefill 写入时标记哪些块属于“首 token 必须”的哪些是“后续可能用到”的。Decode 节点优先拉高优先级块低优先级的可以延迟拉甚至不拉。这个标签在池化层用两个队列实现高优先级队列独立线程处理低优先级队列可以批量合并传输。最后说一个验证方法在 Decode 节点打点记录每个 token 的 KV 等待时间画成热力图。如果热力图显示等待集中在某几层说明那几层的 KV 块特别大或者传输路径有问题。我遇到过某一层的 KV 块因为 head 数不同比其他层大 4 倍单独给那层加预取深度就解决了。分离式推理的调参没有银弹盯着 P99 和热力图比盲目加机器管用。希望帮到你。本文还有配套的精品资源点击获取