九源大模型推理框架 InfiniLM 在国产异构 PD 分离场景取得阶段性进展

近日，InfiniTensor 团队基于九源大模型推理框架 InfiniLM，与 Mooncake 及中科院计算所团队协同开展相关工作，在国产异构 PD 分离推理场景中取得阶段性进展。

随着大模型推理规模持续增长，Prefill-Decode（PD）分离架构正逐渐成为提升系统吞吐、优化资源利用的重要技术方向。近期，InfiniTensor 团队基于 InfiniLM，在国产异构 PD 分离能力建设方面取得了新的进展。

目前，InfiniLM 已支持英伟达、沐曦、摩尔线程、天数智芯等不同计算平台的混合部署，可根据业务需求灵活配置 Prefill 节点与 Decode 节点，构建由多种加速卡协同参与的异构 PD 分离推理集群。系统支持 “x 种加速卡 m 个 P 节点 + y 种加速卡 n 个 D 节点” 的部署模式，并通过 RDMA 通信实现高效的跨节点协作。

目前，相关功能已完成核心验证，正在持续推进性能优化、稳定性验证及工程化整理工作。

打破算力壁垒，让每一份算力物尽其用

随着大模型推理规模不断扩大，Prefill 与 Decode 两个阶段在计算特征上的差异逐渐受到关注。Prefill 阶段以计算密集型负载为主，而 Decode 阶段则更依赖显存访问与数据传输。将两类负载部署在统一资源池中，往往难以充分发挥不同硬件的特点。

针对这一问题，InfiniLM 对 PD（Prefill-Decode）分离架构进行了探索与实践，通过任务解耦和高速互联机制，使 Prefill 与 Decode 节点能够独立部署、协同工作：

Prefill 节点：负责输入 Prompt 的批量计算与 KV Cache 生成；
Decode 节点：基于已生成的 KV Cache 执行逐 Token 推理；
跨节点协同：通过 RDMA 网络传输 KV Cache，降低跨节点数据交换开销，提升异构节点间的协同效率。

核心技术亮点一览

上述进展的背后，是 InfiniLM 在系统架构、调度机制与通信框架等方面的深入探索：

PD 感知的 KV Connector 分层架构：

参考业界先进架构的Scheduler/Worker 分离设计，实现调度逻辑与数据传输的完美解耦。目前已经支持 Mooncake高速传输后端，并预留了接入其他通信后端的接口。

基于 Mooncake 的 RDMA 传输

使用 Mooncake TransferEngine 实现 GPUDirect RDMA，KV Cache 在 P/D 节点 GPU 显存间直传，将传输延迟大幅度降低。

智能两阶段调度与块管理

引入 WAITING_FOR_REMOTE_KVS 新状态，allocate_slots等函数精确协同处理 PD 块分配与 Prefix Cache，缓存命中词元可享“零传输”加速。

高可用异常处理机制

全链路覆盖传输失败、节点异常等场景。实现请求取消的级联通知、KV 传输失败后自动回退重算等机制，确保系统鲁棒性。

从能力验证到持续演进

此次国产异构 PD 分离能力的实现，是 InfiniTensor、Mooncake 与中科院计算所等团队协同推进的阶段性成果。在约两个月的时间里，团队围绕架构设计、通信机制、调度策略与工程实现等方面持续迭代，完成了相关能力的开发与验证。

随着局部优化和代码整理工作的深入，InfiniLM 即将完成多机 RDMA 端到端联调，并推出稳定版本。未来，我们将继续朝着动态扩缩容、跨集群 KV Cache 共享等方向演进，为企业级大模型推理部署提供更强大、更灵活、更高效的基础设施。

我们也期待与更多生态伙伴共同推进国产 AI 推理技术的发展，持续探索异构算力协同与大模型推理优化的更多可能性。