AInfer-PD: In-Place Prefill-Decode Multiplexing for Distributed MoE Agentic RL Rollout

Problem: Prefill-Decode Interleaving in Agentic RL Rollouts

In Agentic RL, models repeatedly generate actions, wait for tool or environment feedback, append new observations, and continue generation. A single rollout contains multiple such episodes, and different trajectories complete tool calls at different times. This causes Prefill and Decode to coexist repeatedly throughout the rollout, rather than handing off once at the start. Prefill batches are large while Decode is latency-sensitive; sharing GPUs leads to interference that slows the entire rollout.

An anonymous internal RL trace with 128 trajectories and 1,265 requests shows 1,137 requests are re-issued in later turns. The median new prompt length is 292 tokens (p99 3,062 tokens), and the median retained context length reaches 4,716 tokens. Later Prefill batches can therefore carry thousands of tokens of history.

Traditional Approach: Disaggregated Prefill and Decode

A common solution splits Prefill and Decode into separate GPU pools, but this requires two sets of compute resources and KV Cache migration between devices.

AInfer-PD: In-Place Multiplexing

Ant Group's AInfer-PD keeps Prefill and Decode on the same GPUs, sharing model weights and KV Cache, while solving two deeper conflicts: cross-rank collective communication ordering and DeepEP mutable communication state contention.

Challenge 1: Cross-Rank Collective Communication Ordering

In large MoE inference, Attention involves both Data Parallelism and Tensor Parallelism, and experts communicate across larger device groups. Prefill and Decode may use intersecting collective communication groups (e.g., Prefill uses full-group TP AllReduce, Decode uses DP-Attention ReduceScatter and AllGather). Different Attention replicas enter Prefill or Decode at different times, causing some ranks to submit Decode before Prefill while others do the opposite. Although each GPU's local stream order is valid, the distributed collectives can form wait cycles that halt progress.

Solution: Rank-Aligned Segment Turnstile

AInfer-PD splits a long Prefill into multiple segments at safe communication boundaries. Each round first aligns all ranks on Decode submission order, then releases one Prefill segment up to the next safe boundary, then yields to the next Decode iteration. This lets a long Prefill span multiple Decode iterations instead of blocking Decode entirely. Crucially, ordering constraints are only applied to collectives that actually conflict; other Prefill computation and non-intersecting communications remain asynchronous.

Challenge 2: DeepEP Communication State Conflicts

DeepEP provides separate communication paths: Normal Path (throughput-oriented, for large Prefill) and Low-Latency Path (for small Decode). These paths share mutable state — buffers, counters, workspaces, events, and QP ranges. Concurrent Prefill and Decode entering different paths corrupt each other's state.

Solution: Separate Mutable Communication States

AInfer-PD keeps model weights, KV storage, physical links, and process-level NVSHMEM runtime shared, but gives Prefill its own Normal Buffer, Counter, Workspace, Stream, and QP Range, while Decode gets independent Low-Latency Buffer, Counter, Workspace, Stream, Event, and QP Range. Isolation is achieved by partitioning only the conflict-prone communication state, not by replicating the entire model.

Experimental Evaluation

Experiments run on NVIDIA H20-3E GPUs using trace replay with the fixed 1,265-request trace.

Single-Node Results (4 GPUs, high Prefill pressure)

vs. AInfer Normal (P/D multiplexing disabled): rollout completion time reduced by 7.1% , 8.7% , 21.5% , and 22.5% across four configurations.

vs. SGLang: completion time reduced by 24.8% to 32.9% .

Ablation: Fine-Grained vs. Global Enqueue

Compared to a Global-Enqueue baseline that only ensures whole-Prefill submission order, the Fine-Grained strategy further reduces completion time by 8.6% to 19.8% . In two configurations, average Decode wait time drops from 127.3 ms and 114.1 ms to 17.5 ms and 18.0 ms .

Dual-Node Results (16 GPUs)

vs. AInfer Normal: rollout completion time reduced by 18.0% to 35.3% .

vs. SGLang: reduced by 18.3% to 31.8% .

Trade-off: dual-node p99 TTFT increases slightly because the optimization target is batch rollout completion, not per-request first-token latency.

Limitations

Validation covers 1–2 nodes of H20-3E; larger clusters and more expert communication runtimes remain untested.

Conclusion

AInfer-PD demonstrates that Agentic RL rollout acceleration is moving beyond single-request inference optimization toward co-design of scheduling, parallelism strategies, and cross-GPU communication.