破解分布式MoE通信冲突！蚂蚁AInfer-PD实现Prefill-Decode就地混跑

智猩猩AI整理

编辑：BugMaker

Agent 做强化学习时，真正耗时间的不只有训练。

在Agentic RL 中，模型往往需要不断经历生成动作、等待工具或环境返回、读取新Observation、继续生成的循环。一次Rollout 可能包含多个这样的回合，而且不同轨迹完成工具调用的时间并不一致。

这会带来一个传统大模型推理里没那么突出的现象。

一些请求还在Decode，另一些请求已经拿到新的环境反馈，又回来重新 Prefill。

Prefill 和 Decode 不再只在推理开始时交接一次，而是会在整轮Agent Rollout 中反复共存。Prefill 批次较大，Decode 又对延迟敏感，两者共用GPU后容易互相干扰，最终拖慢整个 Rollout。

一种常见思路是把Prefill和Decode拆到不同GPU池，但这意味着准备两套计算资源，并在不同设备之间搬运KV Cache。

现在，蚂蚁集团提出了面向分布式MoE Agentic RL Rollout的原地Prefill-Decode复用系统AInfer-PD。

它不拆成两套模型实例，而是让Prefill和Decode继续共享GPU、模型权重和KV Cache，同时解决两类更底层的问题——跨Rank集合通信的执行顺序，以及 DeepEP中Prefill 和Decode共用可变通信状态带来的冲突。

在单节点Prefill高压力实验中，AInfer-PD相比关闭P/D复用的同一AInfer引擎，将固定工作量的 Rollout完成时间降低7.1%到22.5%，相比 SGLang 则降低 24.8% 到 32.9%。

双节点实验中，相比两者分别降低18.0%到35.3%和18.3%到31.8%。

01

Agent多轮交互，让Prefill和Decode

一直撞在一起

普通LLM请求通常先做Prefill，再进入Decode。

Agentic RL却不一样。模型生成一个动作以后，可能需要等待搜索、代码执行、网页访问或者其他环境反馈。新结果回来后，它会被追加到已有上下文中，再执行一次Prefill，然后继续Decode。

不同Agent轨迹又是异步运行的。有人还在生成，有人正在等工具，有人已经带着新Observation 回来了。

所以系统面对的并不是一批Prefill做完以后全部Decode，而是Prefill不断重新插入Decode过程。

论文使用的匿名内部RL Trace包含128条轨迹和1265个请求，其中1137个都是后续轮次重新释放的请求。新增Prompt长度中位数为292 Token，p99达到3062 Token，保留上下文长度中位数已经达到4716 Token。

这意味着后来插进来的Prefill并不一定很小，还可能带着数千甚至上万Token的历史上下文。

真正需要优化的不是某一个请求，而是怎样尽快完成整批相互依赖的Agent轨迹。

过去可以直接把P和D拆到不同设备，但需要独立Worker Pool和KV Cache迁移。AInfer-PD选择另一条路线，模型不复制、KV 不搬走，Prefill和Decode继续在原设备上复用。

02

不拆GPU之后，问题变成了

跨卡通信顺序

让P和D共用GPU，并不是简单增加两条CUDA Stream就结束了。

在大型MoE推理中，Attention往往同时涉及Data Parallelism和Tensor Parallelism，专家部分还会跨更大的设备组通信。

论文指出，Prefill和Decode可能使用彼此相交的Collective通信组。

例如 Prefill走full-group TP AllReduce，Decode则走DP-Attention的 ReduceScatter和AllGather。不同Attention Replica又可能在不同时间进入P 或D。

结果就是一个很隐蔽的问题。

某些Rank先提交D再提交P，另一些Rank却先提交P 再提交D。

每张GPU单独看都没有违反自己的Stream顺序，但放到整个分布式系统里，Collective之间可能形成等待环，最终无法继续推进。

AInfer-PD为此设计了Rank-Aligned Segment Turnstile。

它不再让一个长Prefill从头跑到尾，而是按照安全通信边界拆成多个Segment。每一轮先让所有Rank对Decode的提交顺序达成一致，然后放行一段Prefill，运行到下一个安全边界后，再给下一轮Decode让出机会。

于是一个长Prefill可以跨越多次Decode Iteration，而不是整段时间都把 Decode堵在后面。

更重要的是，AInfer-PD只对真正存在冲突的Collective建立顺序约束，其他 Prefill计算以及不相交的通信仍然可以异步执行。

03

DeepEP也被拆成了两套

独立通信状态

解决Collective顺序还不够，MoE还有另一层冲突来自DeepEP。

DeepEP为不同负载准备了不同通信路径。Normal Path偏向吞吐，更适合大批量Prefill，Low-Latency Path更适合延迟敏感的小批量Decode。

问题是，这两种模式原本更像二选一，并不是按照P和D同时执行来设计的。

它们会共享Buffer、Counter、Workspace、Event和QP等可变状态。一旦Prefill和Decode并发进入不同路径，通信状态就可能互相影响。

AInfer-PD的做法很清楚。

真正可以共享的继续共享，会被不同通信路径修改的状态则分开。

模型权重、KV Storage、物理链路以及进程级NVSHMEM Runtime仍然只有一份。

但Prefill拥有自己的Normal Buffer、Counter、Workspace、Stream 和 QP Range，Decode也拥有独立的Low-Latency Buffer、Counter、Workspace、Stream、Event和 QP Range。

换句话说，AInfer-PD没有用“再复制一套模型”换取隔离，而是把最容易发生冲突的通信状态做成P/D各自所有。

04

单节点较SGLang最多降32.9%，

双节点依旧有效

最终还是要看这些通信设计能不能真正缩短Agent Rollout。

论文在NVIDIA H20-3E GPU上进行了实验，并使用固定的1265个请求进行Trace Replay。

单节点四组Prefill高压力配置中，AInfer-PD相比关闭P/D Multiplexing的 AInfer Normal，将Rollout完成时间分别降低7.1%、8.7%、21.5%和22.5%。

和 SGLang 比较时，完成时间降低24.8%到32.9%。

同引擎消融进一步说明了细粒度通信边界本身的作用。

相比只保证整轮Prefill异步提交顺序的Global-Enqueue，Fine-grained策略将完成时间进一步降低8.6%到19.8%。在两组负载中，Decode平均等待时间也从127.3 ms和114.1 ms降至17.5 ms和18.0 ms。

规模扩到双节点16张H20-3E后，AInfer-PD相比AInfer Normal仍能降低18.0%到35.3%的Rollout完成时间，相比SGLang降低18.3%到31.8%。

这里也存在代价。论文双节点实验中，AInfer-PD的p99 TTFT相比Normal有所增加。作者的优化目标本身就是更快完成整批Rollout，而不是让每一个请求都取得最低首Token延迟。

AInfer-PD的意义并不是再把某一个Kernel做快一点。

它处理的是Agentic RL进入分布式MoE后更系统的问题——当Prefill和Decode在整轮Rollout中持续交错时，怎样继续共享昂贵的模型和KV Cache，同时让跨Rank Collective和DeepEP通信状态不互相打架。

从目前实验范围看，论文主要验证了一到两个节点的H20-3E部署，更大规模集群以及更多Expert Communication Runtime仍未覆盖。

但这项工作已经说明，Agentic RL的Rollout加速正在从单请求推理优化，进一步走向调度、并行策略和跨卡通信协同设计。

关注+星标，获取AI前沿进展与开源一线动态