Source-linked AI summary
psRL: Efficient Training for Agentic AI via Training-Time Prefix Sharing
Mianjie Yu, Zizhao Mo, Huanyu Qu, Zhirong Qian, Huanle Xu, Cen Li, Zifeng Zhao, Zhi Zhou, Jinhua Zhou, Jun Xie, Chengzhong Xu
TL;DR
Structured tree and step-wise RL shift agentic training bottlenecks from rollout to update while producing substantial prefix redundancy. psRL uses prefix-aware scheduling and adaptive KV-cache management, achieving up to 5.2× higher training throughput on production traces.
Problem
Structured sampling increases training sample volume and exposes prefix redundancy that existing training systems largely fail to exploit during updates.
Method
psRL uses global visibility and data immutability to combine flexible prefix-sharing mechanisms with adaptive KV allocation and dynamic caching.
Results
Up to 5.2× improvement in training throughput is achieved on real-world production traces against existing training systems.
Takeaways & Limitations
Prefix-aware scheduling and memory management provide a way to exploit structured agentic-training redundancy while balancing reuse and distributed execution.
Abstract
from arXiv · showhide
In modern agentic AI training, the system bottleneck is shifting from rollout to update. Emerging sampling strategies such as tree-structured and step-wise RL greatly increase training sample volume while incurring relatively low marginal rollout cost, causing the update phase to dominate the end-to-end execution time. Crucially, this shift exposes a new optimization opportunity, as production traces reveal substantial prefix redundancy across training samples. In this paper, we propose psRL (prefix sharing for RL), a new training system for agentic AI designed to exploit prefix redundancy among training samples. Leveraging the global visibility and data immutability inherent to the update phase, psRL achieves efficient workload scheduling and memory management for distributed training. Specifically, psRL introduces two novel prefix-sharing mechanisms that enable flexible, fine-grained workload distribution across GPU workers, simultaneously optimizing prefix reuse and achieving load balancing. Moreover, psRL implements a new underlying KV cache manager that facilitates adaptable block-size allocation and dynamic KV caching, maximizing memory utilization while maintaining a high prefix hit rate. Evaluations using production traces demonstrate that psRL outperforms existing systems by up to 5.2x in throughput. The source code will be publicly available soon.
1 Introduction
Agentic RL is shifting from rollout-bound to update-bound execution as structured sampling increases training samples. psRL exploits shared prefixes through prefix-aware scheduling and adaptive KV-cache management.
- Structured tree and step-wise sampling increase training sample volume while shifting the system bottleneck from rollout to update.These strategies launch parallel exploration from shared prefixes or decompose trajectories into cumulative samples.
- Production traces reveal prefix matching rates frequently exceeding 90% in agentic workloads, creating substantial opportunities to avoid redundant computation.Tree branches share extensive parent prefixes, while step-wise samples recursively include earlier token sequences.
- psRL leverages global visibility and data immutability during training to optimize prefix sharing across the system stack.The entire input dataset is known upfront and remains static within each iteration.
- Inter-batch and self-sequence sharing distribute work flexibly while reusing computed KV states across shared prefixes.Self-sequence sharing partitions long sequences into contiguous chunks assigned to different micro-batches.
- Adaptive block allocation and dynamic block caching improve KV-memory utilization by matching allocation to prefix lengths and caching blocks just in time.The cache manager evicts unreferenced blocks while preserving blocks needed for future operations.
- 5.2× improvement in training throughput is achieved by psRL on real-world production traces against existing training systems.The implementation also accelerates reference-model forward passes and old-policy log-probability computation.
2 Background and Motivation
Structured RL sampling improves sample efficiency and reward density but creates an update-phase bottleneck. Its extensive prefix redundancy motivates training-time prefix sharing, with masking preserving correct computation and gradients.
- Structured RL sampling: Tree-structured and step-wise RL replace independent linear trajectories with shared branches or granular cumulative samples.Tree-step RL combines both structured sampling patterns.
- Bottleneck shift: Structured sampling can inflate total RL time by 1.78×–2.5× over linear trajectory baselines in DTN Agent.The increased latency is attributed to the update phase rather than slower rollout execution.
- Bottleneck shift: Update time becomes dominant under structured sampling, reaching 40%–58% of runtime versus 19%–32% for linear trajectory training.Traditional rollout-dominated workloads allocate 53%–72% of runtime to rollout.
- Prefix redundancy: 94.51% prefix match in ALFWorld Step mode and 90.22% in WebShop Step mode demonstrate substantial structured-sampling redundancy.DTN Agent also reaches 88.92% in Step-wise sampling and 62.24% in Tree-Step sampling.
- Prefix sharing: Training-time prefix sharing reuses common computation while attention masking preserves independent attention results and advantage masking prevents cross-sequence gradient interference.The masks account for sequence-specific attention and distinct advantage values.
3 System Overview
psRL addresses the scheduling and memory-management challenges created by prefix sharing in agentic RL training. It exploits training-time global visibility and data immutability through prefix-sharing mechanisms and coordinated system components.
- Prefix sharing introduces coupled scheduling challenges: systems must increase prefix reuse while balancing workloads across workers and micro-batches.
- Fixed-size KV allocation mismatches variable reusable-prefix lengths, trading memory-access overhead against redundant storage and internal fragmentation.
- psRL leverages global visibility and data immutability to optimize workload scheduling and memory management across the training stack.
- Inter-batch and self-sequence sharing provide flexible, fine-grained workload distribution while reusing prefix KV states.
- A hierarchical distribution strategy combines semantic grouping, worker-level token-wise micro-batching, and prefix-aware execution.
4 Workload Scheduling
psRL schedules semantically related sequences as intact groups, estimates their effective workloads, and balances them across workers before token-wise micro-batching. Its sharing mechanisms preserve reusable prefixes while improving fine-grained distribution.
- Prefix-sharing mechanisms: Inter-batch sharing reuses cached prefix KV states and attention results when identical-prefix sequences span different micro-batches.
- Prefix-sharing mechanisms: Self-sequence sharing partitions long trajectories into contiguous chunks scheduled across micro-batches, reusing prefix KV states for later chunks.
- Semantic grouping: Semantic groups cluster sequences by shared prompts or trajectories, concentrating prefix overlap within groups and limiting cross-group overlap.
- Semantic grouping: psRL keeps semantic groups intact during placement to preserve prefix-sharing opportunities and simplify workload estimation.
- Workload estimation: For each sequence, psRL estimates workload from the uncached suffix using attention, MLP, and system-overhead terms, then sums costs within groups.
- Load balancing: Greedy descending-workload placement assigns intact groups to the least-loaded worker, while token-wise micro-batching limits memory and pipeline imbalance.
5 Memory Management
psRL replaces fixed KV-cache allocation with reuse-aware memory management. It sizes blocks to reusable prefix spans, tracks future references, and evicts blocks using known access patterns.
- Adaptive Block Allocation: Fixed-size blocks either increase access and indirection overhead when small or waste memory and reduce prefix reuse when large.
- Adaptive Block Allocation: Adaptive Block Allocation sizes physical blocks to reusable prefix spans, eliminating redundant prefix materialization while reducing fragmentation and unnecessary accesses.
- Dynamic Block Caching: Known prefix-tree structure and immutable training samples let psRL determine each shared prefix’s span and reuse degree before execution.
- Dynamic Block Caching: During training, reference counts are decremented at micro-batch granularity, and blocks remain in GPU memory only while future reuse exists.
- Dynamic Block Caching: When memory is insufficient, psRL evicts the block with the farthest known reuse distance to minimize recomputation and offloading overhead.
6 System Implementation
psRL is implemented by modifying the distributed RL training stack around prefix-aware orchestration, token-wise execution, adaptive KV caching, and a prefix-sharing runtime. Its mechanisms also extend to other redundant forward-computation stages.
- The implementation is built on veRL with primary modifications to Megatron-LM’s training engine.
- psRL comprises a prefix-aware data orchestrator, token-wise execution planner, adaptive KV cache manager, and prefix-sharing training runtime.
- The same mechanisms extend beyond updates to reference-model forward and old-policy logprobability phases operating on static samples with redundant prefixes.
7 Evaluation
psRL is evaluated across diverse agentic RL workloads and structured sampling modes, with experiments targeting throughput, memory, end-to-end latency, scheduling, memory management, and scalability. It consistently outperforms baseline systems through prefix-aware scheduling and adaptive KV-cache management.
- Experimental Setup: psRL is evaluated on standard Search, WebShop, and ALFWorld agents plus an industrial DTN Agent using five RL strategies.The standard agents use Qwen2.5 models, while the DTN Agent uses Qwen3-235B MoE and supports up to 40K prompt tokens followed by 20K generation tokens per step.
- Main Results: psRL achieves 1.2×–5.2× higher end-to-end training throughput than veRL across four benchmarks.In DTN Agent under Step, psRL reaches 239.1k tokens/s versus 62.7k tokens/s for SGLang-PS, a roughly 3.8× speedup.
- Main Results: psRL maintains memory usage comparable to veRL, while vLLM-PS and SGLang-PS consume 2×–10× more memory than veRL.In ALFWorld Step-S, vLLM-PS peaks at 73.1 GB compared with veRL’s 6.4 GB.
- Main Results: psRL achieves up to 2.1× end-to-end RL speedup over veRL by reducing compute-heavy phases and update latency.For DTN Agent, update latency falls from 462.1 s to 88.7 s; in WebShop Step, psRL completes an iteration in 148.1 s.
- Workload Scheduling: Token-wise micro-batching smooths micro-batch execution, reducing latency spikes and wait-time bubbles while improving global throughput.It repacks and partitions data according to token-level computational density, making forward and backward latency more uniform.
- Memory Management: Dynamic Block Caching reduces DTN Agent peak KV-state memory from 73.0k–73.1k MB to 33.0k MB in Step and Step-RT, exceeding a 54% reduction.Tree-Step memory similarly falls from 72.5k MB to 30.3k MB by reclaiming blocks after final use.
- Scalability: psRL scales near-linearly with cluster size and preserves throughput advantages as sequence length increases from 1× to 8×.On Search Step with 64 devices, psRL reaches 128.5k tokens/s, outperforming veRL by 2.7× and SGLang-PS by 1.9×.
8 Related Works
Related work spans agentic RL systems that either disaggregate or colocate training stages and KV-cache systems developed primarily for online serving. These systems provide context for psRL’s training-focused prefix-sharing design.
- Agentic Systems: Agentic RL systems range from disaggregated rollout-training pipelines to colocated frameworks that overlap multiple RL stages.StreamRL instead uses a one-step off-policy pipeline in which rollout proceeds with slightly stale policy weights.
- KV Management: Serving systems address KV-cache fragmentation and redundant prefix computation through mechanisms such as PagedAttention and RadixAttention.Later work extends KV management to chatbots, streaming, and tool calling.
9 Conclusion
The paper concludes that modern structured RL creates update-phase prefix redundancy that existing training systems largely fail to exploit. psRL uses prefix-aware scheduling and memory management to improve agentic training throughput by up to 5.2×.
- Conclusion: Tree-structured and step-wise RL increase training sample volume at low marginal rollout cost, shifting the bottleneck from rollout to update.These sampling strategies also introduce substantial prefix redundancy among training samples.
- Conclusion: psRL leverages global visibility and data immutability for flexible prefix sharing, load-balanced execution, adaptive KV allocation, and dynamic KV caching.The system is designed specifically for modern agentic RL workloads.
- Conclusion: 5.2× is the maximum agentic training throughput improvement reported for psRL using production traces and real-world workloads.This result summarizes the paper’s extensive evaluation against existing systems.