Search papers, labs, and topics across Lattice.
This paper introduces psRL, a novel training system for agentic AI that capitalizes on prefix redundancy in training samples to enhance efficiency during the update phase. By implementing two innovative prefix-sharing mechanisms and a dynamic KV cache manager, psRL optimizes workload scheduling and memory management across distributed GPU workers. Evaluations show that psRL achieves up to 5.2 times higher throughput compared to existing systems, highlighting its potential to significantly reduce training time in agentic AI applications.
psRL can boost training throughput by over 5x by efficiently sharing prefixes among samples, addressing a critical bottleneck in agentic AI training.
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.
