arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~

AInfer-PD:面向分布式MoE推理的通信安全原地预填-解码复用方案

AInfer-PD: Communication-Safe In-Place Prefill-Decode Multiplexing for Distributed MoE Rollouts

Guowei Wang, Chaokun Yang, Zhenxuan Pan, Yipeng Wei, Yuhong Guo, Minghua Zhu, Zhechuan Zhang, Shuo Wan, Xiaowei Zhu

arXiv 2609.00993首次发表:更新:

AI 中文总结

AInfer-PD是面向分布式MoE推理的通信安全原地预填-解码复用方案,通过协调集体通信顺序与隔离状态,减少预填对解码的干扰,在单节点和双节点场景下均显著降低了推理完成时间。

AI 中文摘要

推理阶段的预填(P)操作通常是大规模强化学习(RL)的主要耗时环节。在智能体RL中,每一条轨迹会在多轮内交替进行模型生成与环境交互,异步轨迹会不断产生新的预填任务,而其他轨迹则处于解码(D)阶段,使得预填与解码的共存成为推理过程的持续特性,而非一次性的提示输入事件。在共享加速器上,这种持续的P/D共存会导致预填干扰对延迟敏感的解码操作,延长推理完成时间。P/D分离方案可避免共置,但需要独立的设备池与KV缓存传输;原地复用方案保留共享设备与KV状态,但现有设计缺乏支持大型MoE部署所需的通信隔离,这类部署结合了张量并行(TP)、数据并行(DP)与分布式专家执行。实际实现中,P和D可能会以不一致的跨秩顺序发出相交的集体通信操作,DeepEP的P和D路径还共享可变协议状态。本文提出AInfer-PD,将原地P/D复用扩展至分布式MoE推理场景,通过跨秩协调P/D的集体通信顺序,为DeepEP的两条路径提供独立的通信状态,使交叉的ADP/ATP与DeepEP路径支持P/D并发执行,同时在同一设备上协调P和D操作,保留共享模型权重与KV存储。在重复的单节点预填密集型工作负载中,与禁用P/D复用的相同AInfer引擎相比,AInfer-PD将固定工作负载的推理完成时间降低7.1%-22.5%,与SGLang相比降低24.8%-32.9%;在双节点场景下,降幅分别为18.0%-35.3%和18.3%-31.8%。在同引擎的 ablation实验中,与全周期异步入队方案相比,细粒度边界设计进一步将完成时间降低8.6%-19.8%。

英文摘要

Rollout inference often dominates the wall-clock time of large-scale reinforcement learning (RL). In agentic RL, each trajectory alternates between model generation and environment interaction over multiple turns. Asynchronous trajectories consequently introduce new prefill (P) work while other trajectories remain in decode (D), making P/D coexistence a persistent property of the rollout rather than a one-time prompt-ingestion event. On shared accelerators, persistent P/D coexistence can make prefill interfere with latency-sensitive decode and prolong rollout completion. P/D disaggregation avoids this co-location but requires separate device pools and KV-cache transfers. In-place multiplexing retains shared devices and KV state, but existing designs lack the communication isolation needed for large MoE deployments that combine attention TP/DP with distributed expert execution. In practical implementations, P and D can issue intersecting collectives in inconsistent cross-rank orders; DeepEP's P and D paths also share mutable protocol state. We present AInfer-PD, which extends in-place P/D multiplexing to distributed MoE rollouts. AInfer-PD coordinates P/D collective order across ranks and gives the two DeepEP paths independent communication state, making crossed ADP/ATP and DeepEP paths safe for concurrent P/D execution. The design retains shared model weights and KV storage while coordinating P and D on the same devices. Across repeated single-node prefill-intensive workloads, AInfer-PD reduces fixed-workload rollout completion time by 7.1-22.5% relative to the same AInfer engine with P/D multiplexing disabled and by 24.8-32.9% relative to SGLang. On two nodes, the reductions are 18.0-35.3% and 18.3-31.8%, respectively. In a same-engine ablation, fine-grained boundaries reduce completion time by a further 8.6-19.8% over whole-epoch asynchronous enqueue.

Comments12 pages, 9 figures

论文原文

arXiv 摘要页 · PDF 原文 · HTML 原文

↑