AMD 推理引擎 ATOM 接上 Mooncake 共享 KV 池：四台 MI350X 服务器上吞吐约为分池两倍
核心概要
编码智能体会把同一段长对话反复送回服务器，历史能否留在显存之外决定了每轮要重算多少 token。AMD 的开源推理引擎 ATOM 把预填充与解码分置在不同 GPU 上，并接入 Mooncake Store，把两侧主机内存合成一个 KV 池，供预填充保存和恢复会话历史。在四台 AMD Instinct MI350X 服务器上，共享池在 96 与 112 并发会话时吞吐约为各服务器独立分池的两倍，KV 复用率保持 94–96%。
深度剖析
智能体编码会话的请求几乎都在重发同一段对话，服务端只要还留着这段历史的 KV，下一轮就只需计算新增的少量 token。 预填充与解码分置在不同 GPU 后，服务器的主机内存也被 NUMA 节点一分为二，解码一侧的那半几乎不存放历史，容量存在却落在预填充够不到的一侧。 单台 8 卡 MI355X 服务器上，ATOM 配合 Mooncake Store 在 96 与 112 并发会话时达到约 44K tokens/s 每 GPU，提示 token 复用率始终高于 95%，接近轨迹允许的 96–97%。
ATOM 把预填充 GPU 显存装不下的历史放进 Mooncake Store，下一轮请求再从池中取回，只计算两层都没留住的 token。 Mooncake Store 把提供内存的服务与使用内存的客户端分开，每个 NUMA 节点各跑一个服务、各贡献 768 GiB DRAM，由 Master 记录对象位置并在池满 90% 时按最近最少使用淘汰。 单台服务器由此得到 1.5 TiB 的 KV 池，约为预填充 GPU 可用 KV 空间的四倍；112 并发会话时，池承担了全部提示 token 的 46%，GPU 显存承担的比例则随会话数增加而下降。
池命中要变成预填充命中，历史必须按预填充的流水线切分对齐地回到四块 GPU 上，缺任何一个阶段都不能算命中。 每个流水级把自己负责的层、按 256 token 一块存成独立对象，Mooncake 把同一块的四个对象放在一个组里一起淘汰，避免只剩部分阶段可加载的残块。 一次典型恢复取回 110K–130K token 历史，每块预填充 GPU 读取并解包的中位时间为 150–220 ms，请求从首个阶段启动到末个阶段完成平均等待 0.4–1.2 s，并随负载上升。
预填充到解码的 KV 交接改成先在预填充侧聚拢、再以大块 RDMA 写入、最后在解码侧散开，网络操作数大幅下降。 预填充按层切分、解码按 token 切分，两者布局不同，直接拷贝会退化成每个 token 每层一次小写入。 一个 8,192 token 块的一个 20 层阶段，直接拷贝需要 163,840 次 576 字节写入，改为聚拢后同一块约 94 MB 只需 12 次写入。
启示与展望
已演示的范围是单台 8 卡 MI355X 服务器与 NeuroSpark 四台 8 卡 MI350X 服务器上的 ATOM 加 Mooncake 部署，负载为 AgentX 回放的编码智能体会话与生产风格重放流量；超过四台服务器的共享池权衡被列为下一步测量对象。
112 并发会话时单机虽仍保持 44K tokens/s 每 GPU，但解码侧 KV 显存中位占用已达 91%，请求开始在预填充前排队，首 token 中位时延升到约 7 s；把池扩展到四台以上服务器会带来更多跨服务器流量和更大的故障域，这一权衡尚未测量。