PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2609.00993v1cs.DC

TL;DR

Agentic RL repeatedly creates prefill while other trajectories decode, and large MoE deployments expose unsafe collective ordering and shared DeepEP state during this coexistence. AInfer-PD extends in-place multiplexing with cross-rank ordering and phase-owned DeepEP state while retaining shared model and KV storage. It reduces rollout completion time across evaluated single- and two-node profiles, with further gains from fine-grained boundaries.

  • Problem

    Large MoE rollouts need communication isolation because intersecting P/D collectives can use inconsistent cross-rank orders and DeepEP’s paths share mutable protocol state.

  • Method

    AInfer-PD coordinates intersecting collective order across ranks and assigns independent mutable communication state to normal-P and low-latency-D DeepEP paths within in-place multiplexing.

  • Results

    AInfer-PD reduces rollout makespan by 7.1–22.5% versus disabled multiplexing and 24.8–32.9% versus SGLang on single-node profiles, with 18.0–35.3% and 18.3–31.8% reductions on two-node profiles.

  • Takeaways & Limitations

    AInfer-PD shortens the rollout critical path without a second model instance or KV transfer, and fine-grained boundaries add 8.6–19.8% over whole-epoch asynchronous enqueue.

  • Takeaways & Limitations

    The evaluation covers one- and two-node H20-3E deployments with BF16/FP8, both MoE backends, and MTP on/off; broader scales and hardware remain open.

Abstract

from arXiv · show

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.

1 Introduction

Agentic RL repeatedly creates prefill while other trajectories decode, so shared accelerators face persistent P/D interference in large MoE rollouts. AInfer-PD adds cross-rank ordering and phase-owned DeepEP state to in-place multiplexing, reducing rollout completion time in evaluated deployments.

  • Motivation: Agentic trajectories repeatedly release continuation prefills while other trajectories decode, making P/D coexistence persistent across multi-turn rollouts.Variable tool latency, generation length, and termination time cause asynchronous continuation prefills.
  • Motivation: Shared accelerators make long P batches compete with latency-sensitive D for compute and communication, increasing rollout completion time.Delaying either phase postpones subsequent trajectory progress.
  • Communication challenges: Large MoE paths can form distributed progress cycles because P and D use intersecting collectives in inconsistent cross-rank orders, while DeepEP paths share mutable protocol state.P may use model-wide TP AllReduce while graph-replayed D uses DP-attention collectives.
  • Contributions: AInfer-PD orders intersecting collectives across ranks and gives normal-P and low-latency-D DeepEP paths phase-owned communication state.The design retains shared model weights and KV storage while supporting concurrent execution.

2 Background and Motivation

Agentic rollouts interleave generation, environment responses, and recurrent prefills, while distributed MoE parallelism exposes intersecting P/D communication paths. Existing disaggregation and in-place approaches leave communication-isolation gaps that AInfer-PD addresses.

  • Agentic rollouts: Each agentic response adds context requiring prefill before generation resumes, so asynchronous trajectories release continuation prefills at different times.The full trajectory set is treated as the unit of rollout work.
  • MoE parallel organization: Large MoE deployments combine attention data parallelism, attention tensor parallelism, and expert communication across wider groups.EP8/ADP2, for example, implies ATP4 when ETP is one.
  • Communication hazards: Intersecting ADP/ATP groups can create distributed progress cycles when P and D enqueue model-wide and DP-attention collectives in different cross-rank orders.Separate P/D communicators do not establish the required inter-communicator order.
  • Communication hazards: DeepEP’s normal and low-latency modes target P and D traffic respectively but share mutable protocol and transport state, preventing concurrent use of both paths.Coexistence therefore falls back to normal mode, slowing D communication relative to the low-latency path.
  • Design motivation: AInfer-PD adds communication isolation for intersecting collectives and concurrent DeepEP modes while retaining shared devices, weights, and KV storage.This extends in-place multiplexing rather than using separate worker pools and KV transfers.

3 Design Overview

AInfer-PD extends in-place multiplexing with rank-shared admission, segmented P execution, cross-rank collective ordering, and phase-owned DeepEP state. These mechanisms interleave successive D iterations with long P executions on shared devices and KV storage.

  • Shared execution: A rank-shared scheduler admits ready D first and uses remaining batch and KV capacity for P on the same model, cache, and GPU pool.Both phases execute against shared weights and KV state.
  • Segmented multiplexing: A long P execution is divided into segments at backend communication boundaries, with each turnstile round enqueuing one D iteration before admitting one P segment.Previously submitted P GPU work continues asynchronously while the current segment runs.
  • Collective ordering: The distributed turnstile establishes identical host-enqueue order across ranks and routes the conflicting P collective after D, while other P work remains asynchronous.Paths without the registered full-group collective retain their original stream-level concurrency.
  • DeepEP isolation: AInfer-PD assigns normal-P and low-latency-D DeepEP operations separate mutable communication state within one process-level runtime.Splitting normal dispatch exposes a boundary before P submits routed data movement, allowing D to proceed.

4 Ordering Phase-Specific Collectives

AInfer-PD removes crossed collective-order cycles by aligning host enqueue and device order at safe P boundaries. A rank-aligned turnstile admits D between P segments without waiting for P GPU completion, while selective routing preserves unrelated overlap.

  • Crossed-group ordering problem: Intersecting P and D collectives can close a global progress cycle even when each rank’s local stream order is valid.Graph-replayed D may reach DP-attention AllGather while P reaches a full-group AllReduce.
  • Ordering protocol: A distributed turnstile uses communication-safe P boundaries to establish a common host-enqueue phase for D, followed by selective stream routing for conflicting device kernels.Only the conflicting pair is ordered; nonintersecting communication keeps background routes.
  • Safe boundaries: P publishes monotone safe boundary identifiers, and the controller advances permitted endpoints by positive segment spans after each post-D phase.The controller rejects stale, duplicate, missing, nonmonotone, and over-permitted boundary events.
  • Segment turnstile: The turnstile’s CPU rendezvous phases align ranks before D enqueue and before releasing the next P segment, without waiting for P GPU completion or collective completion.Submitted P work continues while matched no-work calls keep idle ranks aligned.
  • Selective device ordering: At each crossed-path boundary, D and the following P full-TP AllReduce share a stream, establishing D→P order locally while the turnstile aligns that order across ranks.Repeating the protocol lets D advance through a long P execution; the registered pair itself does not overlap.

5 DeepEP Communication-State Isolation

AInfer-PD isolates DeepEP’s normal-prefill and low-latency-decode communication state within one runtime, allowing both paths to coexist while sharing initialization and physical resources. It also introduces a split-phase boundary that lets decode advance while prefill communication is pending.

  • Communication-state ownership: AInfer-PD adds phase-owned DeepEP communication domains for concurrent normal-P and low-latency-D dispatch.The domains separate buffers, workspaces, counters, events, and queue-pair ranges while retaining one process-level runtime.
  • Communication-state ownership: P and D share NVSHMEM initialization, rank topology, physical links, model weights, and KV storage, but not communication objects modified by kernels.Phase identity selects the appropriate state explicitly, while NIC bandwidth, caches, and unreserved SM capacity remain shared.
  • Safe prefill boundary: AInfer-PD splits normal DeepEP dispatch after notification and before prefill submits data movement.Prefill obtains receive counts in the background and waits for both counts and turnstile permission before submitting data, while decode continues on its low-latency state.
  • Safe prefill boundary: Pending prefill dispatch retains the tensors and stream state needed to finish without blocking decode on its host-side count wait.At most one dispatch remains pending per prefill buffer, and explicit transitions prevent cancellation from reusing partially advanced state.
  • Communication-state ownership: Disjoint per-peer queue-pair intervals and buffer-local synchronization prevent one phase from advancing the other’s transport state.Abandoned device-initiated operations retire their buffers until quiescent cleanup, while shared initialization remains available to either phase or a captured graph.

6 End-to-End Integration

AInfer-PD integrates explicit phase identity, rank-shared admission, KV reservations, and safe-boundary scheduling into the rollout engine. Its segment policy controls the trade-off between decode waiting and coordination overhead without changing the communication contract.

  • Phase-aware integration: Phase identity selects the collective route and DeepEP state from scheduling through attention and expert communication.The identity is explicit rather than inferred from a host thread or CUDA stream.
  • Phase-aware integration: At each scheduling step, the engine commits decode rows, graph state, and optional MTP lookahead before admitting prefill into remaining batch and KV capacity.The rank-shared plan protects decode KV slots and prevents local queues from choosing incompatible phase paths.
  • KV and graph management: Each admitted prefill request reserves KV storage until promotion to decode or cancellation, with validation and draining before reuse.Tensor-parallel workers verify request ownership before promotion, and results from an earlier request lifetime are discarded.
  • KV and graph management: Decode graph replay remains stable while background prefill uses disjoint eager inputs and phase-owned state.MTP drafting and verification form one decode iteration before the post-decode rendezvous releases prefill.
  • Boundary scheduling: The segment policy admits validated boundary intervals in rounds, trading decode waiting time against control overhead under current occupancy and prefill pressure.Fixed spans and a load-selected policy retain the same epoch checks, rendezvous order, and phase-state isolation.

7 Evaluation

The evaluation uses replayed anonymized multi-turn RL workloads across single- and two-node MoE deployments, testing ordering, DeepEP coexistence, segmentation, topology, and rollout pressure. AInfer-PD consistently reduces completion time in prefill-intensive settings, while exposing a completion-time/TTFT trade-off.

  • Methodology: The evaluation replays 1,265 requests from 128 conversations with multi-turn dependencies and reports medians and IQRs across repeated runs.The main workload consumes 6,574,104 input and 220,862 output tokens, with all runs completing without eviction or runtime failure.
  • End-to-end rollout completion: Across H1–H4, AInfer-PD reduces E2E time versus AInfer Normal by 7.1–22.5% and versus SGLang by 24.8–32.9%.The result holds across both evaluated topologies and with MTP enabled or disabled.
  • End-to-end rollout completion: Across H1–H4, p99 request completion falls by 21.3–37.9% versus Normal and 39.3–44.0% versus SGLang.Trajectory completion falls by 8.2–22.7% and 27.0–33.4%, respectively, while p99 TTFT remains within −6.3 to +1.2% of Normal.
  • Collective ordering: Communication-group separation alone stalls the crossed ATP2/ADP2 path, whereas the selective route completes all three full replays.Reducing either topology axis or capping residency at 23 rather than 24 CTAs per SM avoids the stall.
  • Collective ordering: Fine segmentation cuts E2E by 8.6% versus whole-epoch enqueue at the trace-wait point and by a further 19.8% with no added wait.Mean decode wait falls from 127.3/114.1 ms under enqueue to 17.5/18.0 ms under fine segmentation.
  • Topology and DeepEP: With DeepEP and FP8, AInfer-PD reaches 9.18–14.74 requests/s in EP8/ADP8, 26.6–50.4% above Normal and 36.0–62.8% above SGLang.The crossed EP8/ADP2 configuration reaches 9.47–15.74 requests/s across precision and MTP settings.
  • Cross-node and rollout-pressure evaluation: In two-node workloads, AInfer-PD cuts E2E by 18.0–35.3% versus Normal and 18.3–31.8% versus SGLang.Relative to Normal, p99 TTFT rises 13.2–37.9%, exposing the completion-time/TTFT trade-off under completion-oriented admission.
  • Live RL: Two independent live-RL runs reduce summed rollout time from 17.14 to 14.12 hours, a 17.6% decrease.Because generated trajectories may diverge after each step, the result is operational rather than paired evidence.

8 Related Work

Prior work improves rollout-stage coordination, disaggregates or multiplexes prefill and decode, and provides distributed MoE communication substrates. AInfer-PD targets the remaining communication-isolation problem inside a shared distributed MoE rollout worker.

  • Agentic RL execution: RL systems such as HybridFlow, AReaL, and DORA primarily coordinate rollout, reward, and training stages across a cluster.AInfer-PD is complementary because it targets inference-runtime behavior inside a rollout worker.
  • P/D disaggregation: P/D disaggregation separates phase capacity and communication domains but requires distinct workers or pools and KV-state transfers.AInfer-PD instead operates within one in-place model and KV replica.
  • In-place P/D multiplexing: In-place multiplexing systems reduce local P/D contention through granularity changes, stream scheduling, spatial partitioning, or phase-wise resource management.Existing designs establish local P/D overlap but do not cover the intersecting ADP/ATP paths described here.
  • Distributed MoE communication: DeepEP supplies throughput-oriented normal and latency-oriented low-latency protocols for expert dispatch and combine, but its original runtime shares mutable protocol and transport state across modes.AInfer-PD adds phase-owned protocol and transport state while retaining shared process-level initialization and physical resources.

9 Discussion and Limitations

AInfer-PD’s applicability depends on shared-device P/D coexistence and intersecting phase-specific communication paths, with broader hardware and scale left unevaluated. Its state isolation covers mutable protocol state, while scheduling still trades rollout completion against TTFT.

  • Applicability and ordering scope: AInfer-PD targets workers that keep one model and KV state on shared devices while prefill and decode coexist.
  • Applicability and ordering scope: The turnstile applies when phase-specific paths traverse intersecting collective groups in different orders and expose safe insertion boundaries.
  • State isolation and scheduling policy: DeepEP state isolation does not reserve NIC bandwidth, caches, or unreserved SM capacity; admission and communication-SM budgets regulate those resources.
  • State isolation and scheduling policy: Aggressive prefill admission can increase TTFT, and online adaptation of the segment policy remains future work.
  • Evaluation scope: Evaluation covers one- and two-node H20-3E deployments, BF16/FP8, both MoE backends, and MTP on/off; broader scales and hardware remain open.

10 Conclusion

AInfer-PD addresses recurring prefill/decode interference in agentic RL by ordering intersecting collectives and isolating DeepEP phase state while retaining shared model and KV storage. It reduces rollout makespan across evaluated single- and two-node profiles, with further gains from fine-grained segment boundaries.

  • AInfer-PD orders registered collectives and gives DeepEP phase-owned state while retaining one model and KV storage.
  • 7.1–22.5% lower rollout makespan occurs versus the same engine with multiplexing disabled across evaluated prefill-intensive single-node profiles.
  • 24.8–32.9% lower rollout makespan occurs versus SGLang across evaluated prefill-intensive single-node profiles.
  • 8.6–19.8% additional completion improvement comes from fine-grained segment boundaries over whole-epoch asynchronous enqueue in the measured crossed topology.
Loading 2609.00993v1…